Print CUDA state

Author: moelife-coder

Test.py

@@ -24,6 +24,7 @@ def test(data, network):
 network = torch.load(model_name)
 if not args.disable_cuda and torch.cuda.is_available():
     args.device = torch.device('cuda')
+    print("Using CUDA...")
     network.cuda()
 else:
     args.device = torch.device('cpu')

Train.py

@@ -129,6 +129,7 @@ def train(EPOCH, dataloader, test_dataloader, optimizer, loss_function, network,
     network = conv_network.autoencoder(orgsize, bottleneck)
 if not args.disable_cuda and torch.cuda.is_available():
     args.device = torch.device('cuda')
+    print("Using CUDA...")
     network.cuda()
 else:
     args.device = torch.device('cpu')