fix : error device cuda

image_ref/main_ray.py

@@ -88,6 +88,7 @@ def train_model(config,args):
 
             pred_class = torch.argmax(pred_logits, dim=1)
             acc += (pred_class == label).sum().item()
+            print(label.device,pred_logits.device)
             loss = loss_function(pred_logits, label)
             losses += loss.item()
             optimizer.zero_grad()