Mirrors/Track-Anything
1
0
Fork 0
mirror of https://github.com/gaomingqi/Track-Anything.git synced 2025-12-16 16:37:58 +01:00
Code Issues Packages Projects Releases Wiki Activity

xmem load checkpoints cpu -- li

Browse Source
This commit is contained in:
memoryunreal
2023-05-05 07:00:05 +00:00
parent 65588bb237
commit b1ca9273c8
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tracker/model/trainer.py
View File

@@ -195,7 +195,7 @@ class XMemTrainer:
def load_checkpoint(self, path): def load_checkpoint(self, path):
# This method loads everything and should be used to resume training # This method loads everything and should be used to resume training
map_location = 'cuda:%d' % self.local_rank map_location = 'cuda:%d' % self.local_rank
checkpoint = torch.load(path, map_location={'cuda:0': map_location}) checkpoint = torch.load(path, map_location={'cpu': map_location})
it = checkpoint['it'] it = checkpoint['it']
network = checkpoint['network'] network = checkpoint['network']
@@ -218,7 +218,7 @@ class XMemTrainer:
def load_network(self, path): def load_network(self, path):
# This method loads only the network weight and should be used to load a pretrained model # This method loads only the network weight and should be used to load a pretrained model
map_location = 'cuda:%d' % self.local_rank map_location = 'cuda:%d' % self.local_rank
src_dict = torch.load(path, map_location={'cuda:0': map_location}) src_dict = torch.load(path, map_location={'cpu': map_location})
self.load_network_in_memory(src_dict) self.load_network_in_memory(src_dict)
print(f'Network weight loaded from {path}') print(f'Network weight loaded from {path}')