minor

pomonam · Jul 10, 2024 · fd679d2 · fd679d2
1 parent e528662
commit fd679d2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/kronfluence/module/tracked_module.py b/kronfluence/module/tracked_module.py
@@ -97,7 +97,7 @@ def __init__(
         self._constant: torch.Tensor = nn.Parameter(
             torch.zeros(
                 1,
-                dtype=self.original_module.weight.dtype,
+                dtype=torch.bfloat16,
                 requires_grad=True,
             )
         )