diff --git a/image_ref/main_ray.py b/image_ref/main_ray.py
index 2cb7285312d8223b6cc7f8c4826baa1b5aa33193..3530a8b46cfb67dd8d0f13d18a449e07218e9a48 100644
--- a/image_ref/main_ray.py
+++ b/image_ref/main_ray.py
@@ -54,6 +54,7 @@ def train_model(config,args):
     n_class = len(data_train.dataset.classes)
     weight = torch.Tensor([1/n_class,1-1/n_class])
     weight.to(device)
+    print('weight',weight.device)
     loss_function = nn.CrossEntropyLoss(weight=weight)
 
     # Load existing checkpoint through `get_checkpoint()` API.