Christina Theodoris
commited on
Commit
·
f953554
1
Parent(s):
efc403d
fix to properly move model after checking device
Browse files
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 ==
|
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 |
|