From 61852dac3047ee71322a4a34a0d20325f308c22b Mon Sep 17 00:00:00 2001 From: Juhan Bae Date: Wed, 13 Mar 2024 02:20:58 -0400 Subject: [PATCH] Check contiguous tensor --- kronfluence/module/tracked_module.py | 1 + 1 file changed, 1 insertion(+) diff --git a/kronfluence/module/tracked_module.py b/kronfluence/module/tracked_module.py index 06ec347..88c6af2 100644 --- a/kronfluence/module/tracked_module.py +++ b/kronfluence/module/tracked_module.py @@ -592,6 +592,7 @@ def backward_hook(output_gradient: torch.Tensor) -> None: storage=self._storage, damping=self.score_args.damping, ) + assert preconditioned_gradient.is_contiguous() self._cached_per_sample_gradient = None preconditioned_gradient = preconditioned_gradient.to(dtype=self.score_args.score_dtype)