mlrun 1.6.4rc2__py3-none-any.whl → 1.7.0rc20__py3-none-any.whl

mlrun/frameworks/pytorch/callbacks/tensorboard_logging_callback.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 from datetime import datetime
-from typing import Callable, Dict, List, Tuple, Union
+from typing import Callable, Union
 
 import torch
 from torch import Tensor
@@ -63,7 +63,7 @@ class _PyTorchTensorboardLogger(TensorboardLogger):
 
     def __init__(
         self,
-        statistics_functions: List[
+        statistics_functions: list[
             Callable[[Union[Parameter]], Union[float, Parameter]]
         ],
         context: mlrun.MLClientCtx = None,
@@ -94,7 +94,7 @@ class _PyTorchTensorboardLogger(TensorboardLogger):
                                       update. Notice that writing to tensorboard too frequently may cause the training
                                       to be slower. Default: 'epoch'.
         """
-        super(_PyTorchTensorboardLogger, self).__init__(
+        super().__init__(
             statistics_functions=statistics_functions,
             context=context,
             tensorboard_directory=tensorboard_directory,
@@ -249,19 +249,19 @@ class TensorboardLoggingCallback(LoggingCallback):
         context: mlrun.MLClientCtx = None,
         tensorboard_directory: str = None,
         run_name: str = None,
-        weights: Union[bool, List[str]] = False,
-        statistics_functions: List[
+        weights: Union[bool, list[str]] = False,
+        statistics_functions: list[
             Callable[[Union[Parameter, Tensor]], Union[float, Tensor]]
         ] = None,
-        dynamic_hyperparameters: Dict[
+        dynamic_hyperparameters: dict[
             str,
-            Tuple[
+            tuple[
                 str,
-                Union[List[Union[str, int]], Callable[[], PyTorchTypes.TrackableType]],
+                Union[list[Union[str, int]], Callable[[], PyTorchTypes.TrackableType]],
             ],
         ] = None,
-        static_hyperparameters: Dict[
-            str, Union[PyTorchTypes.TrackableType, Tuple[str, List[Union[str, int]]]]
+        static_hyperparameters: dict[
+            str, Union[PyTorchTypes.TrackableType, tuple[str, list[Union[str, int]]]]
         ] = None,
         update_frequency: Union[int, str] = "epoch",
         auto_log: bool = False,
@@ -322,7 +322,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         :raise MLRunInvalidArgumentError: In case both 'context' and 'tensorboard_directory' parameters were not given
                                           or the 'update_frequency' was incorrect.
         """
-        super(TensorboardLoggingCallback, self).__init__(
+        super().__init__(
             dynamic_hyperparameters=dynamic_hyperparameters,
             static_hyperparameters=static_hyperparameters,
             auto_log=auto_log,
@@ -345,7 +345,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         # Save the configurations:
         self._tracked_weights = weights
 
-    def get_weights(self) -> Dict[str, Parameter]:
+    def get_weights(self) -> dict[str, Parameter]:
         """
         Get the weights tensors tracked. The weights will be stored in a dictionary where each key is the weight's name
         and the value is the weight's parameter (tensor).
@@ -354,7 +354,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         """
         return self._logger.weights
 
-    def get_weights_statistics(self) -> Dict[str, Dict[str, List[float]]]:
+    def get_weights_statistics(self) -> dict[str, dict[str, list[float]]]:
         """
         Get the weights mean results logged. The results will be stored in a dictionary where each key is the weight's
         name and the value is a list of mean values per epoch.
@@ -365,7 +365,7 @@ class TensorboardLoggingCallback(LoggingCallback):
 
     @staticmethod
     def get_default_weight_statistics_list() -> (
-        List[Callable[[Union[Parameter, Tensor]], Union[float, Tensor]]]
+        list[Callable[[Union[Parameter, Tensor]], Union[float, Tensor]]]
     ):
         """
         Get the default list of statistics functions being applied on the tracked weights each epoch.
@@ -381,7 +381,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         validation_set: DataLoader = None,
         loss_function: Module = None,
         optimizer: Optimizer = None,
-        metric_functions: List[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
         scheduler=None,
     ):
         """
@@ -396,7 +396,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         :param metric_functions: The metric functions to be stored in this callback.
         :param scheduler:        The scheduler to be stored in this callback.
         """
-        super(TensorboardLoggingCallback, self).on_setup(
+        super().on_setup(
             model=model,
             training_set=training_set,
             validation_set=validation_set,
@@ -439,7 +439,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         for logging. Epoch 0 (pre-run state) will be logged here.
         """
         # Setup all the results and hyperparameters dictionaries:
-        super(TensorboardLoggingCallback, self).on_run_begin()
+        super().on_run_begin()
 
         # Log the initial summary of the run:
         self._logger.write_initial_summary_text()
@@ -470,10 +470,10 @@ class TensorboardLoggingCallback(LoggingCallback):
         # Write the final summary of the run:
         self._logger.write_final_summary_text()
 
-        super(TensorboardLoggingCallback, self).on_run_end()
+        super().on_run_end()
 
     def on_validation_end(
-        self, loss_value: PyTorchTypes.MetricValueType, metric_values: List[float]
+        self, loss_value: PyTorchTypes.MetricValueType, metric_values: list[float]
     ):
         """
         Before the validation (in a training case it will be per epoch) ends, this method will be called to log the
@@ -482,9 +482,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         :param loss_value:    The loss summary of this validation.
         :param metric_values: The metrics summaries of this validation.
         """
-        super(TensorboardLoggingCallback, self).on_validation_end(
-            loss_value=loss_value, metric_values=metric_values
-        )
+        super().on_validation_end(loss_value=loss_value, metric_values=metric_values)
 
         # Check if this run was part of an evaluation:
         if not self._is_training:
@@ -503,7 +501,7 @@ class TensorboardLoggingCallback(LoggingCallback):
 
         :param epoch: The epoch that has just ended.
         """
-        super(TensorboardLoggingCallback, self).on_epoch_end(epoch=epoch)
+        super().on_epoch_end(epoch=epoch)
 
         # Log the weights statistics:
         self._logger.log_weights_statistics()
@@ -540,9 +538,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         :param y_true: The true value part of the current batch.
         :param y_pred: The prediction (output) of the model for this batch's input ('x').
         """
-        super(TensorboardLoggingCallback, self).on_train_batch_end(
-            batch=batch, x=x, y_true=y_true, y_pred=y_pred
-        )
+        super().on_train_batch_end(batch=batch, x=x, y_true=y_true, y_pred=y_pred)
 
         # Write the batch loss and metrics results to their graphs:
         self._logger.write_training_results()
@@ -559,9 +555,7 @@ class TensorboardLoggingCallback(LoggingCallback):
         :param y_true: The true value part of the current batch.
         :param y_pred: The prediction (output) of the model for this batch's input ('x').
         """
-        super(TensorboardLoggingCallback, self).on_validation_batch_end(
-            batch=batch, x=x, y_true=y_true, y_pred=y_pred
-        )
+        super().on_validation_batch_end(batch=batch, x=x, y_true=y_true, y_pred=y_pred)
 
         # Write the batch loss and metrics results to their graphs:
         self._logger.write_validation_results()

mlrun/frameworks/pytorch/callbacks_handler.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from typing import Dict, List, Tuple, Union
+from typing import Union
 
 from torch import Tensor
 from torch.nn import Module
@@ -66,7 +66,7 @@ class CallbacksHandler:
     A class for handling multiple callbacks during a run.
     """
 
-    def __init__(self, callbacks: List[Union[Callback, Tuple[str, Callback]]]):
+    def __init__(self, callbacks: list[Union[Callback, tuple[str, Callback]]]):
         """
         Initialize the callbacks handler with the given callbacks he will handle. The callbacks can be passed as their
         initialized instances or as a tuple where [0] is a name that will be attached to him and [1] will be the
@@ -99,7 +99,7 @@ class CallbacksHandler:
                 self._callbacks[callback.__class__.__name__] = callback
 
     @property
-    def callbacks(self) -> Dict[str, Callback]:
+    def callbacks(self) -> dict[str, Callback]:
         """
         Get the callbacks dictionary handled by this handler.
 
@@ -114,9 +114,9 @@ class CallbacksHandler:
         validation_set: DataLoader,
         loss_function: Module,
         optimizer: Optimizer,
-        metric_functions: List[PyTorchTypes.MetricFunctionType],
+        metric_functions: list[PyTorchTypes.MetricFunctionType],
         scheduler,
-        callbacks: List[str] = None,
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_setup' method of every callback in the callbacks list. If the list is 'None' (not given), all
@@ -145,7 +145,7 @@ class CallbacksHandler:
             scheduler=scheduler,
         )
 
-    def on_run_begin(self, callbacks: List[str] = None) -> bool:
+    def on_run_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_run_begin' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -159,7 +159,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_run_end(self, callbacks: List[str] = None) -> bool:
+    def on_run_end(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_run_end' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -173,7 +173,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_epoch_begin(self, epoch: int, callbacks: List[str] = None) -> bool:
+    def on_epoch_begin(self, epoch: int, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_epoch_begin' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -189,7 +189,7 @@ class CallbacksHandler:
             epoch=epoch,
         )
 
-    def on_epoch_end(self, epoch: int, callbacks: List[str] = None) -> bool:
+    def on_epoch_end(self, epoch: int, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_epoch_end' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -205,7 +205,7 @@ class CallbacksHandler:
             epoch=epoch,
         )
 
-    def on_train_begin(self, callbacks: List[str] = None) -> bool:
+    def on_train_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_train_begin' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -219,7 +219,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_train_end(self, callbacks: List[str] = None) -> bool:
+    def on_train_end(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_train_end' method of every callback in the callbacks list. If the list is 'None' (not given), all
         callbacks will be called.
@@ -233,7 +233,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_validation_begin(self, callbacks: List[str] = None) -> bool:
+    def on_validation_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_validation_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -250,8 +250,8 @@ class CallbacksHandler:
     def on_validation_end(
         self,
         loss_value: PyTorchTypes.MetricValueType,
-        metric_values: List[float],
-        callbacks: List[str] = None,
+        metric_values: list[float],
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_validation_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -271,7 +271,7 @@ class CallbacksHandler:
         )
 
     def on_train_batch_begin(
-        self, batch: int, x, y_true: Tensor, callbacks: List[str] = None
+        self, batch: int, x, y_true: Tensor, callbacks: list[str] = None
     ) -> bool:
         """
         Call the 'on_train_batch_begin' method of every callback in the callbacks list. If the list is 'None'
@@ -298,7 +298,7 @@ class CallbacksHandler:
         x,
         y_pred: Tensor,
         y_true: Tensor,
-        callbacks: List[str] = None,
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_train_batch_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -322,7 +322,7 @@ class CallbacksHandler:
         )
 
     def on_validation_batch_begin(
-        self, batch: int, x, y_true: Tensor, callbacks: List[str] = None
+        self, batch: int, x, y_true: Tensor, callbacks: list[str] = None
     ) -> bool:
         """
         Call the 'on_validation_batch_begin' method of every callback in the callbacks list. If the list is 'None'
@@ -349,7 +349,7 @@ class CallbacksHandler:
         x,
         y_pred: Tensor,
         y_true: Tensor,
-        callbacks: List[str] = None,
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_validation_batch_end' method of every callback in the callbacks list. If the list is 'None'
@@ -375,7 +375,7 @@ class CallbacksHandler:
     def on_inference_begin(
         self,
         x,
-        callbacks: List[str] = None,
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_inference_begin' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -396,7 +396,7 @@ class CallbacksHandler:
         self,
         y_pred: Tensor,
         y_true: Tensor,
-        callbacks: List[str] = None,
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_inference_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -415,7 +415,7 @@ class CallbacksHandler:
             y_true=y_true,
         )
 
-    def on_train_loss_begin(self, callbacks: List[str] = None) -> bool:
+    def on_train_loss_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_train_loss_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -430,7 +430,7 @@ class CallbacksHandler:
         )
 
     def on_train_loss_end(
-        self, loss_value: PyTorchTypes.MetricValueType, callbacks: List[str] = None
+        self, loss_value: PyTorchTypes.MetricValueType, callbacks: list[str] = None
     ) -> bool:
         """
         Call the 'on_train_loss_end' method of every callback in the callbacks list. If the list is 'None' (not given),
@@ -447,7 +447,7 @@ class CallbacksHandler:
             loss_value=loss_value,
         )
 
-    def on_validation_loss_begin(self, callbacks: List[str] = None) -> bool:
+    def on_validation_loss_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_validation_loss_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -462,7 +462,7 @@ class CallbacksHandler:
         )
 
     def on_validation_loss_end(
-        self, loss_value: PyTorchTypes.MetricValueType, callbacks: List[str] = None
+        self, loss_value: PyTorchTypes.MetricValueType, callbacks: list[str] = None
     ) -> bool:
         """
         Call the 'on_validation_loss_end' method of every callback in the callbacks list. If the list is 'None'
@@ -479,7 +479,7 @@ class CallbacksHandler:
             loss_value=loss_value,
         )
 
-    def on_train_metrics_begin(self, callbacks: List[str] = None) -> bool:
+    def on_train_metrics_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_train_metrics_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -495,8 +495,8 @@ class CallbacksHandler:
 
     def on_train_metrics_end(
         self,
-        metric_values: List[PyTorchTypes.MetricValueType],
-        callbacks: List[str] = None,
+        metric_values: list[PyTorchTypes.MetricValueType],
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_train_metrics_end' method of every callback in the callbacks list. If the list is 'None'
@@ -513,7 +513,7 @@ class CallbacksHandler:
             metric_values=metric_values,
         )
 
-    def on_validation_metrics_begin(self, callbacks: List[str] = None) -> bool:
+    def on_validation_metrics_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_validation_metrics_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -529,8 +529,8 @@ class CallbacksHandler:
 
     def on_validation_metrics_end(
         self,
-        metric_values: List[PyTorchTypes.MetricValueType],
-        callbacks: List[str] = None,
+        metric_values: list[PyTorchTypes.MetricValueType],
+        callbacks: list[str] = None,
     ) -> bool:
         """
         Call the 'on_validation_metrics_end' method of every callback in the callbacks list. If the list is 'None'
@@ -547,7 +547,7 @@ class CallbacksHandler:
             metric_values=metric_values,
         )
 
-    def on_backward_begin(self, callbacks: List[str] = None) -> bool:
+    def on_backward_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_backward_begin' method of every callback in the callbacks list. If the list is 'None' (not given),
         all callbacks will be called.
@@ -561,7 +561,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_backward_end(self, callbacks: List[str] = None) -> bool:
+    def on_backward_end(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_backward_end' method of every callback in the callbacks list. If the list is 'None' (not given),
         all callbacks will be called.
@@ -575,7 +575,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_optimizer_step_begin(self, callbacks: List[str] = None) -> bool:
+    def on_optimizer_step_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_optimizer_step_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -589,7 +589,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_optimizer_step_end(self, callbacks: List[str] = None) -> bool:
+    def on_optimizer_step_end(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_optimizer_step_end' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -603,7 +603,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_scheduler_step_begin(self, callbacks: List[str] = None) -> bool:
+    def on_scheduler_step_begin(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_scheduler_step_begin' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -617,7 +617,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def on_scheduler_step_end(self, callbacks: List[str] = None) -> bool:
+    def on_scheduler_step_end(self, callbacks: list[str] = None) -> bool:
         """
         Call the 'on_scheduler_step_end' method of every callback in the callbacks list. If the list is 'None'
         (not given), all callbacks will be called.
@@ -631,7 +631,7 @@ class CallbacksHandler:
             callbacks=self._parse_names(names=callbacks),
         )
 
-    def _parse_names(self, names: Union[List[str], None]) -> List[str]:
+    def _parse_names(self, names: Union[list[str], None]) -> list[str]:
         """
         Parse the given callbacks names. If they are not 'None' then the names will be returned as they are, otherwise
         all of the callbacks handled by this handler will be returned (the default behavior of when there were no names
@@ -646,7 +646,7 @@ class CallbacksHandler:
         return list(self._callbacks.keys())
 
     def _run_callbacks(
-        self, method_name: str, callbacks: List[str], *args, **kwargs
+        self, method_name: str, callbacks: list[str], *args, **kwargs
     ) -> bool:
         """
         Run the given method from the 'CallbackInterface' on all the specified callbacks with the given arguments.

mlrun/frameworks/pytorch/mlrun_interface.py

@@ -14,7 +14,7 @@
 #
 import importlib
 import sys
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any, Union
 
 import torch
 import torch.multiprocessing as mp
@@ -109,13 +109,13 @@ class PyTorchMLRunInterface:
         loss_function: Module,
         optimizer: Optimizer,
         validation_set: DataLoader = None,
-        metric_functions: List[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
         scheduler=None,
         scheduler_step_frequency: Union[int, float, str] = "epoch",
         epochs: int = 1,
         training_iterations: int = None,
         validation_iterations: int = None,
-        callbacks: List[Callback] = None,
+        callbacks: list[Callback] = None,
         use_cuda: bool = True,
         use_horovod: bool = None,
     ):
@@ -221,12 +221,12 @@ class PyTorchMLRunInterface:
         self,
         dataset: DataLoader,
         loss_function: Module = None,
-        metric_functions: List[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
         iterations: int = None,
-        callbacks: List[Callback] = None,
+        callbacks: list[Callback] = None,
         use_cuda: bool = True,
         use_horovod: bool = None,
-    ) -> List[PyTorchTypes.MetricValueType]:
+    ) -> list[PyTorchTypes.MetricValueType]:
         """
         Initiate an evaluation process on this interface configuration.
 
@@ -303,9 +303,9 @@ class PyTorchMLRunInterface:
     def add_auto_logging_callbacks(
         self,
         add_mlrun_logger: bool = True,
-        mlrun_callback_kwargs: Dict[str, Any] = None,
+        mlrun_callback_kwargs: dict[str, Any] = None,
         add_tensorboard_logger: bool = True,
-        tensorboard_callback_kwargs: Dict[str, Any] = None,
+        tensorboard_callback_kwargs: dict[str, Any] = None,
     ):
         """
         Get automatic logging callbacks to both MLRun's context and Tensorboard. For further features of logging to both
@@ -347,7 +347,7 @@ class PyTorchMLRunInterface:
 
     def predict(
         self,
-        inputs: Union[Tensor, List[Tensor]],
+        inputs: Union[Tensor, list[Tensor]],
         use_cuda: bool = True,
         batch_size: int = -1,
     ) -> Tensor:
@@ -402,13 +402,13 @@ class PyTorchMLRunInterface:
         loss_function: Module = None,
         optimizer: Optimizer = None,
         validation_set: DataLoader = None,
-        metric_functions: List[PyTorchTypes.MetricFunctionType] = None,
+        metric_functions: list[PyTorchTypes.MetricFunctionType] = None,
         scheduler=None,
         scheduler_step_frequency: Union[int, float, str] = "epoch",
         epochs: int = 1,
         training_iterations: int = None,
         validation_iterations: int = None,
-        callbacks: List[Callback] = None,
+        callbacks: list[Callback] = None,
         use_cuda: bool = True,
         use_horovod: bool = None,
     ):
@@ -734,7 +734,7 @@ class PyTorchMLRunInterface:
 
     def _validate(
         self, is_evaluation: bool = False
-    ) -> Tuple[PyTorchTypes.MetricValueType, List[PyTorchTypes.MetricValueType]]:
+    ) -> tuple[PyTorchTypes.MetricValueType, list[PyTorchTypes.MetricValueType]]:
         """
         Initiate a single epoch validation.
 
@@ -817,7 +817,7 @@ class PyTorchMLRunInterface:
         )
         return loss_value, metric_values
 
-    def _print_results(self, loss_value: Tensor, metric_values: List[float]):
+    def _print_results(self, loss_value: Tensor, metric_values: list[float]):
         """
         Print the given result between each epoch.
 
@@ -832,7 +832,7 @@ class PyTorchMLRunInterface:
             + tabulate(table, headers=["Metrics", "Values"], tablefmt="pretty")
         )
 
-    def _metrics(self, y_pred: Tensor, y_true: Tensor) -> List[float]:
+    def _metrics(self, y_pred: Tensor, y_true: Tensor) -> list[float]:
         """
         Call all the metrics on the given batch's truth and prediction output.
 
@@ -860,7 +860,7 @@ class PyTorchMLRunInterface:
         average_tensor = self._hvd.allreduce(rank_value, name=name)
         return average_tensor.item()
 
-    def _get_learning_rate(self) -> Union[Tuple[str, List[Union[str, int]]], None]:
+    def _get_learning_rate(self) -> Union[tuple[str, list[Union[str, int]]], None]:
         """
         Try and get the learning rate value form the stored optimizer.
 
@@ -949,8 +949,8 @@ class PyTorchMLRunInterface:
 
     @staticmethod
     def _tensor_to_cuda(
-        tensor: Union[Tensor, Dict, List, Tuple],
-    ) -> Union[Tensor, Dict, List, Tuple]:
+        tensor: Union[Tensor, dict, list, tuple],
+    ) -> Union[Tensor, dict, list, tuple]:
         """
         Send to given tensor to cuda if it is a tensor. If the given object is a dictionary, the dictionary values will
         be sent to the function again recursively. If the given object is a list or a tuple, all the values in it will
@@ -997,7 +997,7 @@ class PyTorchMLRunInterface:
         dataset: DataLoader,
         iterations: int,
         description: str,
-        metrics: List[PyTorchTypes.MetricFunctionType],
+        metrics: list[PyTorchTypes.MetricFunctionType],
     ) -> tqdm:
         """
         Create a progress bar for training and validating / evaluating.
@@ -1028,8 +1028,8 @@ class PyTorchMLRunInterface:
     @staticmethod
     def _update_progress_bar(
         progress_bar: tqdm,
-        metrics: List[PyTorchTypes.MetricFunctionType],
-        values: List[PyTorchTypes.MetricValueType],
+        metrics: list[PyTorchTypes.MetricFunctionType],
+        values: list[PyTorchTypes.MetricValueType],
     ):
         """
         Update the progress bar metrics results.