rezav
2 years ago
1 changed files with 3 additions and 0 deletions
+ 3
- 0
predictor/cross_validation.py
View File
| def cv(args, out_dir): | 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')) | save_args(args, os.path.join(out_dir, 'args.json')) | ||||
| test_loss_file = os.path.join(out_dir, 'test_loss.pkl') | test_loss_file = os.path.join(out_dir, 'test_loss.pkl') | ||||
| print("cuda available: " + str(torch.cuda.is_available())) | print("cuda available: " + str(torch.cuda.is_available())) |
Loading…