bugfix

mokawi · Mar 29, 2021 · db75a4c · db75a4c
1 parent 95b7534
commit db75a4c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/clctm.py b/clctm.py
@@ -154,10 +154,11 @@ def __init__(
 
         if use_cuda is None:
             self.use_cuda = torch.cuda.is_available()
-            self.cvmodel = self.cvmodel.to(f'cuda:{cuda_device}')
         else:
             self.use_cuda = use_cuda
-        self.cuda_device = cuda_device
+        self.cuda_device = cuda_device if cuda_device is not None else 
+        if self.use_cuda:
+            self.cvmodel = self.cvmodel.to(f'cuda:{self.cuda_device}')
 
         if data is not None:
             self.add_docs(data)