enable mmap in torch.load (#2479)

wenet-e2e · Apr 15, 2024 · b1dadcb · b1dadcb
1 parent c9bc0aa
commit b1dadcb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/wenet/utils/checkpoint.py b/wenet/utils/checkpoint.py
@@ -25,7 +25,7 @@
 
 def load_checkpoint(model: torch.nn.Module, path: str) -> dict:
     logging.info('Checkpoint: loading from checkpoint %s' % path)
-    checkpoint = torch.load(path, map_location='cpu')
+    checkpoint = torch.load(path, map_location='cpu', mmap=True)
     missing_keys, unexpected_keys = model.load_state_dict(checkpoint,
                                                           strict=False)
     for key in missing_keys: