fix : error device cuda

image_ref/main_ray.py

@@ -53,7 +53,8 @@ def train_model(config,args):
     # init training
     n_class = len(data_train.dataset.classes)
     weight = torch.Tensor([1/n_class,1-1/n_class])
-    weight.to(device)
+    if torch.cuda.is_available():
+        weight = weight.cuda()
     print('weight',weight.device)
     loss_function = nn.CrossEntropyLoss(weight=weight)