Updated traker.py to include a lambda_reg term in arguments

trak/traker.py

@@ -64,6 +64,8 @@ def __init__(
         proj_max_batch_size: int = 32,
         projector_seed: int = 0,
         grad_wrt: Optional[Iterable[str]] = None,
+        lambda_reg: int = 0
+
     ) -> None:
         """
 
@@ -129,6 +131,9 @@ def __init__(
                 as they appear in the model's state dictionary. If None,
                 gradients are taken with respect to all model parameters.
                 Defaults to None.
+            lambda_reg  (int):
+                Applies L2 regularization to the xtx term in scorecomputers
+                with form xtx + lambda_reg*1. Defaults to 0
 
         """
 
@@ -138,6 +143,7 @@ def __init__(
         self.device = device
         self.dtype = ch.float16 if use_half_precision else ch.float32
         self.grad_wrt = grad_wrt
+        self.lambda_reg = lambda_reg
 
         logging.basicConfig()
         self.logger = logging.getLogger("TRAK")
@@ -183,7 +189,7 @@ def __init__(
         if score_computer is None:
             score_computer = BasicScoreComputer
         self.score_computer = score_computer(
-            dtype=self.dtype, device=self.device, logging_level=logging_level
+            dtype=self.dtype, device=self.device, logging_level=logging_level, lambda_reg = self.lambda_reg
         )
 
         metadata = {