SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

adding option to load on cpu

Browse Source
This commit is contained in:
thomwolf
2019-02-08 10:37:40 +01:00
parent 777459b471
commit 5ee4f17234
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
pytorch_pretrained_bert/modeling.py
View File

@@ -580,7 +580,7 @@ class BertPreTrainedModel(nn.Module):
model = cls(config, *inputs, **kwargs) model = cls(config, *inputs, **kwargs)
if state_dict is None and not from_tf: if state_dict is None and not from_tf:
weights_path = os.path.join(serialization_dir, WEIGHTS_NAME) weights_path = os.path.join(serialization_dir, WEIGHTS_NAME)
state_dict = torch.load(weights_path) state_dict = torch.load(weights_path, map_location='cpu' if not torch.cuda.is_available() else None)
if tempdir: if tempdir:
# Clean up temp dir # Clean up temp dir
shutil.rmtree(tempdir) shutil.rmtree(tempdir)