Skip to content
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

Jump to bottom

UniTS model hard-coded "cuda:0" when loading pretrain weights #23

Open
zhangzw16 opened this issue Jul 16, 2024 · 2 comments
Open

UniTS model hard-coded "cuda:0" when loading pretrain weights #23

zhangzw16 opened this issue Jul 16, 2024 · 2 comments
Assignees
@zhangzw16
Labels
bug Something isn't working

Comments

@zhangzw16
Copy link
Collaborator

ProbTS/probts/model/forecaster/point_forecaster/units.py

Line 1014 in d299e1c

ckpt = torch.load(pretrain_weight_path, map_location='cuda:0')

This should be made more robust to support various devices.
@zhangzw16 zhangzw16 added the bug Something isn't working label Jul 16, 2024
@zhangzw16 zhangzw16 self-assigned this Jul 16, 2024
@zhangzw16
Copy link
Collaborator Author

Fixed with #24

@zhangzw16 zhangzw16 reopened this Jul 21, 2024
@zhangzw16
Copy link
Collaborator Author

Maybe need a better way to specify device when inference.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@zhangzw16 zhangzw16

Labels
bug Something isn't working
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
1 participant
@zhangzw16