Update utils.py

ambitious-octopus · Sep 12, 2024 · d1c7cba · d1c7cba
1 parent 1b3854f
commit d1c7cba
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/model_compression_toolkit/core/pytorch/utils.py b/model_compression_toolkit/core/pytorch/utils.py
@@ -38,10 +38,7 @@ def set_model(model: torch.nn.Module, train_mode: bool = False):
         model.eval()
 
     device = get_working_device()
-    if device.type == CPU:
-        model.cpu()
-    else:
-        model.cuda()
+    model.to(device)
 
 
 def to_torch_tensor(tensor,
@@ -112,4 +109,4 @@ def clip_inf_values_float16(tensor: Tensor) -> Tensor:
     # Replace inf values with max float16 value
     tensor[inf_mask] = MAX_FLOAT16 * torch.sign(tensor[inf_mask])
 
-    return tensor
+    return tensor