Make bias optional for MatMul nodes (#1637)

neuralmagic · Jun 23, 2023 · 0b64b31 · 0b64b31
1 parent a6fdb7f
commit 0b64b31
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/sparseml/yolov8/trainers.py b/src/sparseml/yolov8/trainers.py
@@ -555,7 +555,7 @@ def __init__(self, model="yolov8n.yaml", type="v8") -> None:
 
         if model_str.endswith(".pt"):
             if os.path.exists(model_str):
-                ckpt = torch.load(model_str)
+                ckpt = torch.load(model_str, map_location="cpu")
                 self.is_sparseml_checkpoint = (
                     "source" in ckpt and ckpt["source"] == "sparseml"
                 )