Try to empty cache before exiting

Author: moelife-coder

Train.py

@@ -150,4 +150,5 @@ def train(EPOCH, dataloader, test_dataloader, optimizer, loss_function, network,
 test_loader = data.DataLoader(dataset = test_data, batch_size = BATCH_SIZE, shuffle = False)
 optimizer = torch.optim.Adam(network.parameters(), lr = LR)
 train(EPOCH, train_loader, test_loader, optimizer, loss_func, network, model_name, negative_loss, args.device, should_view, f)
-f.close()
+f.close()
+torch.cuda.empty_cache()