Christina Theodoris commited on
Commit
f953554
·
1 Parent(s): efc403d

fix to properly move model after checking device

Browse files
Files changed (1) hide show
  1. geneformer/perturber_utils.py +1 -1
geneformer/perturber_utils.py CHANGED
@@ -213,7 +213,7 @@ def move_to_cuda(model):
213
  # get what device model is currently on
214
  model_device = next(model.parameters()).device
215
  # Check if the model is on the CPU and move to cuda if necessary
216
- if (model_device.type == 'cpu') and (device == "cuda"):
217
  model.to(device)
218
 
219
 
 
213
  # get what device model is currently on
214
  model_device = next(model.parameters()).device
215
  # Check if the model is on the CPU and move to cuda if necessary
216
+ if (model_device.type == "cpu") and (device.type == "cuda"):
217
  model.to(device)
218
 
219