@@ -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())) |