Xiao Zeqi commited on
Update exp_base.py
Browse files- experiments/exp_base.py +1 -0
experiments/exp_base.py
CHANGED
|
@@ -56,6 +56,7 @@ def load_custom_checkpoint(algo, checkpoint_path):
|
|
| 56 |
model_path = hf_hub_download(repo_id=repo_id, filename=file_name)
|
| 57 |
ckpt = torch.load(model_path, map_location=torch.device('cpu'))
|
| 58 |
|
|
|
|
| 59 |
filtered_state_dict = {}
|
| 60 |
for k, v in ckpt['state_dict'].items():
|
| 61 |
if "frame_timestep_embedder" in k:
|
|
|
|
| 56 |
model_path = hf_hub_download(repo_id=repo_id, filename=file_name)
|
| 57 |
ckpt = torch.load(model_path, map_location=torch.device('cpu'))
|
| 58 |
|
| 59 |
+
# quick workaround
|
| 60 |
filtered_state_dict = {}
|
| 61 |
for k, v in ckpt['state_dict'].items():
|
| 62 |
if "frame_timestep_embedder" in k:
|