Skip to content

GitLab

"git@developer.sourcefind.cn:OpenDAS/ollama.git" did not exist on "29a8975c66860c2b0b3fa15515d3e3dd4c78d51d"
Unverified Commit 50f69e70 authored by Jiamin's avatar Jiamin Committed by GitHub
Browse files

fix: _get_max_memory when model has no attr output_device (#416)

parent 71e32749
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
mmcv/runner/hooks/logger/text.py
View file @ 50f69e70
...@@ -50,7 +50,7 @@ class TextLoggerHook(LoggerHook): ...@@ -50,7 +50,7 @@ class TextLoggerHook(LoggerHook):
self._dump_log(runner.meta, runner) self._dump_log(runner.meta, runner)
def _get_max_memory(self, runner): def _get_max_memory(self, runner):
device = runner.model.output_device device = getattr(runner.model, 'output_device', None)
mem = torch.cuda.max_memory_allocated(device=device) mem = torch.cuda.max_memory_allocated(device=device)
mem_mb = torch.tensor([mem / (1024 * 1024)], mem_mb = torch.tensor([mem / (1024 * 1024)],
dtype=torch.int, dtype=torch.int,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment