fix : error device cuda

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.