Fix load_model to transfer full_tensor

Tharmeekan · Tharmeekan · commit 152cc6dd6c05 · 2026-05-26T11:11:52.000+02:00
diff --git a/src/weathergen/model/model_interface.py b/src/weathergen/model/model_interface.py
@@ -204,7 +204,7 @@ def load_model(cf, model, device, run_id: str, mini_epoch=-1):
                 )
                 maybe_sharded_sd[param_name] = torch.nn.Parameter(sharded_tensor)
             else:
-                maybe_sharded_sd[param_name] = full_tensor
+                maybe_sharded_sd[param_name] = full_tensor.to(device)
         # choose `assign=True` for sharded model since we cannot call `copy_` on meta tensor
         mkeys, ukeys = model.load_state_dict(maybe_sharded_sd, strict=False, assign=True)
 

Original file line number	Diff line number	Diff line change
`@@ -204,7 +204,7 @@ def load_model(cf, model, device, run_id: str, mini_epoch=-1):`
`204`	`204`	`)`
`205`	`205`	`maybe_sharded_sd[param_name] = torch.nn.Parameter(sharded_tensor)`
`206`	`206`	`else:`
`207`		`- maybe_sharded_sd[param_name] = full_tensor`
	`207`	`+ maybe_sharded_sd[param_name] = full_tensor.to(device)`
`208`	`208`	# choose `assign=True` for sharded model since we cannot call `copy_` on meta tensor
`209`	`209`	`mkeys, ukeys = model.load_state_dict(maybe_sharded_sd, strict=False, assign=True)`
`210`	`210`