memory leak

main_ray_tune.py

@@ -199,9 +199,7 @@ def test_best_model(best_result, args):
     if torch.cuda.is_available():
         device = "cuda:0"
         if torch.cuda.device_count() > 1:
-            print(type(best_trained_model))
             best_trained_model = torch.nn.DataParallel(best_trained_model)
-            print(type(best_trained_model))
 
     best_trained_model.to(device)
     criterion_rt = torch.nn.MSELoss()