Fix loading without cuda

marble.py

@@ -59,7 +59,10 @@ def setup_control_mlp(
 
     net = control_mlp(features)
     net.load_state_dict(
-        torch.load(os.path.join(file_dir, f"model_weights/{material_parameter}.pt"))
+        torch.load(
+            os.path.join(file_dir, f"model_weights/{material_parameter}.pt"),
+            map_location=device
+        )
     )
     net.to(device, dtype=dtype)
     net.eval()