Commit 1b7cb35b authored by Artem Pulkin's avatar Artem Pulkin
Browse files

workflows: print torch device as well

parent 93786234
Pipeline #86018 passed with stages
in 24 minutes and 33 seconds
......@@ -1296,7 +1296,7 @@ class FitWorkflow(Workflow):
if torch.is_tensor(obj):
size = obj.nelement() * obj.element_size()
if size > threshold:
self.log.debug(f" tensor {obj.shape} {obj.dtype} {pretty_memory(size)}")
self.log.debug(f" tensor {obj.shape} {obj.dtype} {obj.device} {pretty_memory(size)}")
except Exception as e:" failed to evaluate {obj} with exception {e}")
if next(self.nn[0].parameters()).is_cuda:
Supports Markdown
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