Higher preicison for grad computation

pomonam · Jul 14, 2024 · 78adca8 · 78adca8
1 parent 4772696
commit 78adca8
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/examples/openwebtext/compute_scores.py b/examples/openwebtext/compute_scores.py
@@ -96,6 +96,7 @@ def main():
     # We can invest some time in getting more accurate SVD results.
     score_args.use_full_svd = True
     score_args.precondition_dtype = torch.float32
+    score_args.per_sample_gradient_dtype = torch.float32
     analyzer.compute_pairwise_scores(
         scores_name=args.scores_name,
         score_args=score_args,