Spaces:
Running
Running
Update skyreels_v2_infer/modules/t5.py
Browse files
skyreels_v2_infer/modules/t5.py
CHANGED
|
@@ -434,7 +434,7 @@ class T5EncoderModel(ModelMixin):
|
|
| 434 |
# init model
|
| 435 |
model = umt5_xxl(encoder_only=True, return_tokenizer=False)
|
| 436 |
logging.info(f"loading {checkpoint_path}")
|
| 437 |
-
model.load_state_dict(torch.load(checkpoint_path, map_location="cpu"))
|
| 438 |
self.model = model
|
| 439 |
if shard_fn is not None:
|
| 440 |
self.model = shard_fn(self.model, sync_module_states=False)
|
|
|
|
| 434 |
# init model
|
| 435 |
model = umt5_xxl(encoder_only=True, return_tokenizer=False)
|
| 436 |
logging.info(f"loading {checkpoint_path}")
|
| 437 |
+
model.load_state_dict(torch.load(checkpoint_path, map_location="cpu", weights_only=True))
|
| 438 |
self.model = model
|
| 439 |
if shard_fn is not None:
|
| 440 |
self.model = shard_fn(self.model, sync_module_states=False)
|