memory ):

predictor/cross_validation.py

@@ -104,6 +104,9 @@ def create_model(data, hidden_size, gpu_id=None):
 
 
 def cv(args, out_dir):
+    torch.cuda.set_per_process_memory_fraction(0.6, 0)
+    # Clear any cached memory
+    torch.cuda.empty_cache()
     save_args(args, os.path.join(out_dir, 'args.json'))
     test_loss_file = os.path.join(out_dir, 'test_loss.pkl')
     print("cuda available: " + str(torch.cuda.is_available()))