Skip to content

Ensure checkpoint loading works with PyTorch 2.6 #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yeonsungk1m wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Ensure checkpoint loading works with PyTorch 2.6 #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion in_the_wild/videopose_diffusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,11 @@ def main(args):
chk_file_path = "./checkpoint/in_the_wild_best_epoch.bin"
# chk_filename = args.resume or args.evaluate
print('Loading checkpoint', chk_file_path)
checkpoint = torch.load(chk_file_path, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_file_path,
map_location=lambda storage, loc: storage,
weights_only=False,
)
print('This model was trained for {} epochs'.format(checkpoint['epoch']))
model_pos.load_state_dict(checkpoint['model_pos'], strict=False)

Expand Down
18 changes: 15 additions & 3 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -304,7 +304,11 @@ def fetch(subjects, action_filter=None, subset=1, parse_3d_poses=True):
chk_filename = os.path.join(args.checkpoint, args.resume if args.resume else args.evaluate)
# chk_filename = args.resume or args.evaluate
print('Loading checkpoint', chk_filename)
checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_filename,
map_location=lambda storage, loc: storage,
weights_only=False,
)
Comment on lines 304 to +311

@chatgpt-codex-connector chatgpt-codex-connector Bot Oct 26, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Guard torch.load weights_only for pre-2.6 PyTorch

The new torch.load(..., weights_only=False) calls assume the argument exists, but earlier PyTorch releases (e.g. 1.x) do not accept weights_only and raise TypeError: torch.load() got an unexpected keyword argument 'weights_only' before any checkpoint is loaded. The README still advertises support for PyTorch ≥0.4, so this change will crash on previously supported environments. Consider only passing weights_only when running on PyTorch versions that provide the parameter (e.g. via feature detection or try/except). The same pattern is repeated in the other files modified in this commit.

Useful? React with 👍 / 👎.

print('This model was trained for {} epochs'.format(checkpoint['epoch']))
model_pos_train.load_state_dict(checkpoint['model_pos'], strict=False)
model_pos.load_state_dict(checkpoint['model_pos'], strict=False)
Expand Down Expand Up @@ -890,7 +894,11 @@ def evaluate(test_generator, action=None, return_predictions=False, use_trajecto
model_eval = newmodel
chk_file_path = '/mnt/data3/home/zjl/workspace/3dpose/PoseFormer/checkpoint/train_pf_00/epoch_60.bin'
print('Loading evaluate checkpoint of comparison model', chk_file_path)
checkpoint = torch.load(chk_file_path, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_file_path,
map_location=lambda storage, loc: storage,
weights_only=False,
)
model_eval.load_state_dict(checkpoint['model_pos'], strict=False)
model_eval.eval()
else:
Expand All @@ -903,7 +911,11 @@ def evaluate(test_generator, action=None, return_predictions=False, use_trajecto
elif args.evaluate != '':
chk_file_path = os.path.join(args.checkpoint, args.evaluate)
print('Loading evaluate checkpoint', chk_file_path)
checkpoint = torch.load(chk_file_path, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_file_path,
map_location=lambda storage, loc: storage,
weights_only=False,
)
print('This model was trained for {} epochs'.format(checkpoint['epoch']))
# model_pos_train.load_state_dict(checkpoint['model_pos'], strict=False)
model_eval.load_state_dict(checkpoint['model_pos'])
Expand Down
18 changes: 15 additions & 3 deletions main_3dhp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -411,7 +411,11 @@ def fetch(subjects, action_filter=None, subset=1, parse_3d_poses=True):
chk_filename = os.path.join(args.checkpoint, args.resume if args.resume else args.evaluate)
# chk_filename = args.resume or args.evaluate
print('Loading checkpoint', chk_filename)
checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_filename,
map_location=lambda storage, loc: storage,
weights_only=False,
)
print('This model was trained for {} epochs'.format(checkpoint['epoch']))
model_pos_train.load_state_dict(checkpoint['model_pos'], strict=False)
model_pos.load_state_dict(checkpoint['model_pos'], strict=False)
Expand Down Expand Up @@ -1062,7 +1066,11 @@ def evaluate(test_generator, action=None, return_predictions=False, use_trajecto
model_eval = newmodel
chk_file_path = '/mnt/data3/home/zjl/workspace/3dpose/PoseFormer/checkpoint/train_pf_00/epoch_60.bin'
print('Loading evaluate checkpoint of comparison model', chk_file_path)
checkpoint = torch.load(chk_file_path, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_file_path,
map_location=lambda storage, loc: storage,
weights_only=False,
)
model_eval.load_state_dict(checkpoint['model_pos'], strict=False)
model_eval.eval()
else:
Expand All @@ -1075,7 +1083,11 @@ def evaluate(test_generator, action=None, return_predictions=False, use_trajecto
elif args.evaluate != '':
chk_file_path = os.path.join(args.checkpoint, args.evaluate)
print('Loading evaluate checkpoint', chk_file_path)
checkpoint = torch.load(chk_file_path, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_file_path,
map_location=lambda storage, loc: storage,
weights_only=False,
)
print('This model was trained for {} epochs'.format(checkpoint['epoch']))
# model_pos_train.load_state_dict(checkpoint['model_pos'], strict=False)
model_eval.load_state_dict(checkpoint['model_pos'])
Expand Down
18 changes: 15 additions & 3 deletions main_draw.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -269,7 +269,11 @@ def fetch(subjects, action_filter=None, subset=1, parse_3d_poses=True):
chk_filename = os.path.join(args.checkpoint, args.resume if args.resume else args.evaluate)
# chk_filename = args.resume or args.evaluate
print('Loading checkpoint', chk_filename)
checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_filename,
map_location=lambda storage, loc: storage,
weights_only=False,
)
print('This model was trained for {} epochs'.format(checkpoint['epoch']))
model_pos_train.load_state_dict(checkpoint['model_pos'], strict=False)
model_pos.load_state_dict(checkpoint['model_pos'], strict=False)
Expand Down Expand Up @@ -340,7 +344,11 @@ def evaluate(test_generator, action=None, return_predictions=False, use_trajecto
model_eval = newmodel
chk_file_path = '/mnt/data3/home/zjl/workspace/3dpose/PoseFormer/checkpoint/train_pf_00/epoch_60.bin'
print('Loading evaluate checkpoint of comparison model', chk_file_path)
checkpoint = torch.load(chk_file_path, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_file_path,
map_location=lambda storage, loc: storage,
weights_only=False,
)
model_eval.load_state_dict(checkpoint['model_pos'], strict=False)
model_eval.eval()
else:
Expand All @@ -353,7 +361,11 @@ def evaluate(test_generator, action=None, return_predictions=False, use_trajecto
elif args.evaluate != '':
chk_file_path = os.path.join(args.checkpoint, args.evaluate)
print('Loading evaluate checkpoint', chk_file_path)
checkpoint = torch.load(chk_file_path, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_file_path,
map_location=lambda storage, loc: storage,
weights_only=False,
)
print('This model was trained for {} epochs'.format(checkpoint['epoch']))
# model_pos_train.load_state_dict(checkpoint['model_pos'], strict=False)
model_eval.load_state_dict(checkpoint['model_pos'])
Expand Down
18 changes: 15 additions & 3 deletions main_in_the_wild.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,7 +264,11 @@ def fetch(subjects, action_filter=None, subset=1, parse_3d_poses=True):
chk_filename = os.path.join(args.checkpoint, args.resume if args.resume else args.evaluate)
# chk_filename = args.resume or args.evaluate
print('Loading checkpoint', chk_filename)
checkpoint = torch.load(chk_filename, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_filename,
map_location=lambda storage, loc: storage,
weights_only=False,
)
print('This model was trained for {} epochs'.format(checkpoint['epoch']))
model_pos_train.load_state_dict(checkpoint['model_pos'], strict=False)
model_pos.load_state_dict(checkpoint['model_pos'], strict=False)
Expand Down Expand Up @@ -624,7 +628,11 @@ def evaluate(test_generator, action=None, return_predictions=False, use_trajecto
model_eval = newmodel
chk_file_path = '/mnt/data3/home/zjl/workspace/3dpose/PoseFormer/checkpoint/train_pf_00/epoch_60.bin'
print('Loading evaluate checkpoint of comparison model', chk_file_path)
checkpoint = torch.load(chk_file_path, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_file_path,
map_location=lambda storage, loc: storage,
weights_only=False,
)
model_eval.load_state_dict(checkpoint['model_pos'], strict=False)
model_eval.eval()
else:
Expand All @@ -637,7 +645,11 @@ def evaluate(test_generator, action=None, return_predictions=False, use_trajecto
elif args.evaluate != '':
chk_file_path = os.path.join(args.checkpoint, args.evaluate)
print('Loading evaluate checkpoint', chk_file_path)
checkpoint = torch.load(chk_file_path, map_location=lambda storage, loc: storage)
checkpoint = torch.load(
chk_file_path,
map_location=lambda storage, loc: storage,
weights_only=False,
)
print('This model was trained for {} epochs'.format(checkpoint['epoch']))
# model_pos_train.load_state_dict(checkpoint['model_pos'], strict=False)
model_eval.load_state_dict(checkpoint['model_pos'])
Expand Down