minor

pomonam · pomonam · commit 2c2d765bd3d2 · 2024-07-10T18:14:00.000-04:00
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=torch.bfloat16,
+                dtype=torch.float16,
                 requires_grad=True,
             )
         )

Original file line number	Diff line number	Diff line change
`@@ -97,7 +97,7 @@ def __init__(`
`97`	`97`	`self._constant: torch.Tensor = nn.Parameter(`
`98`	`98`	`torch.zeros(`
`99`	`99`	`1,`
`100`		`- dtype=torch.bfloat16,`
	`100`	`+ dtype=torch.float16,`
`101`	`101`	`requires_grad=True,`
`102`	`102`	`)`
`103`	`103`	`)`