Skip to content

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

Open
@zhangzw16

Description

@zhangzw16

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

This should be made more robust to support various devices.

Metadata

Metadata

Assignees

Labels

bugSomething isn't working

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions