make model cuda if available

baoy-nlp · Apr 6, 2018 · e1df0eb · e1df0eb
1 parent 4342595
commit e1df0eb
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/inference.py b/inference.py
@@ -35,6 +35,9 @@ def main(args):
     model.load_state_dict(torch.load(args.load_checkpoint))
     print("Model loaded from %s"%(args.load_checkpoint))
 
+    if torch.cuda.is_available():
+        model = model.cuda()
+
     model.eval()
 
     samples, z = model.inference(n=args.num_samples)