Skip to content
Snippets Groups Projects
Unverified Commit e1c7674b authored by Mashiro's avatar Mashiro Committed by GitHub
Browse files

[Fix] Fix potential OOM error when loading checkpoint with `map_location="cuda:0"` (#665)

* [Fix] Change the default value of map_location in PretrainedInit to cpu

* Fix as comment
parent 618a063d
No related branches found
No related tags found
No related merge requests found
...@@ -469,10 +469,10 @@ class PretrainedInit: ...@@ -469,10 +469,10 @@ class PretrainedInit:
initialize. For example, if we would like to only load the initialize. For example, if we would like to only load the
backbone of a detector model, we can set ``prefix='backbone.'``. backbone of a detector model, we can set ``prefix='backbone.'``.
Defaults to None. Defaults to None.
map_location (str): map tensors into proper locations. map_location (str): map tensors into proper locations. Defaults to cpu.
""" """
def __init__(self, checkpoint, prefix=None, map_location=None): def __init__(self, checkpoint, prefix=None, map_location='cpu'):
self.checkpoint = checkpoint self.checkpoint = checkpoint
self.prefix = prefix self.prefix = prefix
self.map_location = map_location self.map_location = map_location
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment