Callback fixes

allennlp/training/callbacks/log_writer.py

@@ -291,13 +291,13 @@ def log_epoch(
     def _should_log_distributions_next_batch(self) -> bool:
         return (
             self._distribution_interval is not None
-            and (self.trainer._batch_num_total + 1) % self._distribution_interval == 0  # type: ignore[union-attr]
+            and (self.trainer._total_batches_completed + 1) % self._distribution_interval == 0
         )
 
     def _should_log_distributions_this_batch(self) -> bool:
         return (
             self._distribution_interval is not None
-            and self.trainer._batch_num_total % self._distribution_interval == 0  # type: ignore[union-attr]
+            and self.trainer._total_batches_completed % self._distribution_interval == 0
         )
 
     def _enable_activation_logging(self) -> None:
@@ -318,7 +318,7 @@ def hook(module_, inputs, outputs):
                 self._module_hook_handles.append(module.register_forward_hook(hook))
 
     def _should_log_this_batch(self) -> bool:
-        return self.trainer._batch_num_total % self._summary_interval == 0  # type: ignore[union-attr]
+        return self.trainer._total_batches_completed % self._summary_interval == 0  # type: ignore[union-attr]
 
     def _log_activation_distribution(self, outputs: Any, module_name: str) -> None:
         activations_to_log: Dict[str, torch.Tensor] = {}

allennlp/training/callbacks/tensorboard.py

@@ -49,7 +49,7 @@ def log_scalars(
         log_prefix: str = "",
         epoch: Optional[int] = None,
     ) -> None:
-        timestep = epoch if epoch is not None else self.trainer._batch_num_total  # type: ignore[union-attr]
+        timestep = epoch if epoch is not None else self.trainer._total_batches_completed  # type: ignore[union-attr]
         log = self._train_log if not log_prefix.startswith("validation") else self._validation_log
         for key, value in scalars.items():
             name = f"{log_prefix}/{key}" if log_prefix else key
@@ -59,7 +59,7 @@ def log_scalars(
     def log_tensors(
         self, tensors: Dict[str, torch.Tensor], log_prefix: str = "", epoch: Optional[int] = None
     ) -> None:
-        timestep = epoch if epoch is not None else self.trainer._batch_num_total  # type: ignore[union-attr]
+        timestep = epoch if epoch is not None else self.trainer._total_batches_completed  # type: ignore[union-attr]
         log = self._train_log if not log_prefix.startswith("validation") else self._validation_log
         for key, values in tensors.items():
             name = f"{log_prefix}/{key}" if log_prefix else key

allennlp/training/callbacks/wandb.py

@@ -134,7 +134,7 @@ def _log(
             dict_to_log = {f"{log_prefix}/{k}": v for k, v in dict_to_log.items()}
         if epoch is not None:
             dict_to_log["epoch"] = epoch
-        self.wandb.log(dict_to_log, step=self.trainer._batch_num_total)  # type: ignore[union-attr]
+        self.wandb.log(dict_to_log, step=self.trainer._total_batches_completed)  # type: ignore[union-attr]
 
     @overrides
     def on_start(