remove : debugging print

image_ref/main_ray.py

@@ -56,7 +56,6 @@ def train_model(config,args):
     weight.float()
     if torch.cuda.is_available():
         weight = weight.cuda()
-    print('weight',weight.device)
     loss_function = nn.CrossEntropyLoss(weight=weight)
 
     # Load existing checkpoint through `get_checkpoint()` API.
@@ -94,7 +93,6 @@ 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()
@@ -195,9 +193,7 @@ def test_model(best_result, args):
     if torch.cuda.is_available():
         device = "cuda:0"
         if torch.cuda.device_count() > 1:
-            print(type(model))
             net = torch.nn.DataParallel(model)
-            print(type(net))
     model.to(device)
     # init training
     loss_function = nn.CrossEntropyLoss()