From 57442ed4939be330a7912bf5548155b5317bda49 Mon Sep 17 00:00:00 2001 From: Juhan Bae Date: Fri, 5 Jul 2024 02:14:03 -0400 Subject: [PATCH] Print device --- kronfluence/computer/factor_computer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/kronfluence/computer/factor_computer.py b/kronfluence/computer/factor_computer.py index db8a76d..051fa1a 100644 --- a/kronfluence/computer/factor_computer.py +++ b/kronfluence/computer/factor_computer.py @@ -307,7 +307,7 @@ def fit_covariance_matrices( if torch.is_tensor(obj) or ( hasattr(obj, 'data') and torch.is_tensor(obj.data)) and obj.device == torch.device( "cuda"): - print(type(obj), obj.size()) + print(type(obj), obj.size(), obj.device) except: pass @@ -345,7 +345,7 @@ def fit_covariance_matrices( if torch.is_tensor(obj) or ( hasattr(obj, 'data') and torch.is_tensor(obj.data)) and obj.device == torch.device( "cuda"): - print(type(obj), obj.size()) + print(type(obj), obj.size(), obj.device) except: pass