>>101399676
diff --git a/chameleon/inference/loader.py b/chameleon/inference/loader.py
index 167d70f..e898d08 100644
--- a/chameleon/inference/loader.py
+++ b/chameleon/inference/loader.py
@@ -20,7 +20,7 @@ def _convert(model_args: ModelArgs, consolidated_path: Path) -> Transformer:
model = Transformer(model_args)
transfer_results = model.load_state_dict(
- torch.load(str(consolidated_path), map_location='cuda'),
+ torch.load(str(consolidated_path)),
strict=False,
)
Changing that made it load at least with a 24GB GPU.