This commit is contained in:
Patrick von Platen 2020-06-11 21:28:01 +02:00 коммит произвёл GitHub
Родитель 6293eb04df
Коммит caf3746678
Не найден ключ, соответствующий данной подписи
Идентификатор ключа GPG: 4AEE18F83AFDEB23
1 изменённых файлов: 5 добавлений и 6 удалений

Просмотреть файл

@ -104,8 +104,8 @@ class PyTorchBenchmark(Benchmark):
)
torch.cuda.reset_max_memory_cached()
# calculate loss and do backpropagation
_train()
# calculate loss and do backpropagation
_train()
elif not self.args.no_tpu and is_torch_tpu_available():
# tpu
raise NotImplementedError(
@ -129,8 +129,7 @@ class PyTorchBenchmark(Benchmark):
logger.info(
"Please consider updating PyTorch to version 1.4 to get more accuracy on GPU memory usage"
)
memory = Memory(torch.cuda.max_memory_cached())
memory = Memory(torch.cuda.max_memory_reserved())
memory = Memory(torch.cuda.max_memory_reserved())
return memory, summary
else:
@ -215,8 +214,8 @@ class PyTorchBenchmark(Benchmark):
)
torch.cuda.reset_max_memory_cached()
# run forward
_forward()
# run forward
_forward()
elif not self.args.no_tpu and is_torch_tpu_available():
# tpu
raise NotImplementedError(