add spacing for twa sampling

Author: sapkotaruz11

dicee/weight_averaging.py

@@ -395,7 +395,7 @@ class TWA(AbstractCallback):
 
     def __init__(self, twa_start_epoch: int, lr_init: float,
                  num_samples: int = 5, reg_lambda: float = 0.0,
-                 max_epochs: int = None):
+                 max_epochs: int = None, twa_c_epochs :int = 5):
         """
         Parameters
         ----------
@@ -409,13 +409,16 @@ def __init__(self, twa_start_epoch: int, lr_init: float,
             Regularization coefficient for β updates.
         max_epochs : int
             Total number of training epochs.
+        twa_c_epochs : int
+            Spacing (in epochs) between consecutive weight samples for TWA.
         """
         super().__init__()
         self.twa_start_epoch = twa_start_epoch
         self.num_samples = num_samples
         self.reg_lambda = reg_lambda
         self.max_epochs = max_epochs
         self.lr_init = lr_init
+        self.twa_c_epochs = twa_c_epochs
 
         # State variables
         self.current_epoch = -1
@@ -464,7 +467,7 @@ def on_train_epoch_start(self, trainer, model):
     def on_train_epoch_end(self, trainer, model):
         """Main TWA logic: build subspace and update in β space."""
         # Step 1: collect weight samples before TWA starts
-        if self.current_epoch < self.twa_start_epoch:
+        if self.current_epoch < self.twa_start_epoch and self.current_epoch  % self.twa_c_epochs == 0 :
             self.sample_weights(model)  # rolling buffer handled inside
             return