Adopt type hinting generics in standard collections (PEP 585)

1 files changed, 5 insertions, 5 deletions

diff --git a/models/model.py b/models/model.py
index 369d6c2..cb0e756 100644
--- a/models/model.py
+++ b/models/model.py
@@ -1,4 +1,4 @@
-from typing import List, Dict, Union, Tuple
+from typing import Union
 
 import torch
 from torch.utils.data.dataloader import default_collate
@@ -7,15 +7,15 @@ from torch.utils.data.dataloader import default_collate
 class Model:
     def __init__(
             self,
-            batch_size: Tuple[int, int]
+            batch_size: tuple[int, int]
     ):
         (self.pr, self.k) = batch_size
 
     def _batch_splitter(
             self,
-            batch: List[Dict[str, Union[str, torch.Tensor]]]
-    ) -> List[Tuple[Dict[str, List[Union[str, torch.Tensor]]],
-                    Dict[str, List[Union[str, torch.Tensor]]]]]:
+            batch: list[dict[str, Union[str, torch.Tensor]]]
+    ) -> list[tuple[dict[str, list[Union[str, torch.Tensor]]],
+                    dict[str, list[Union[str, torch.Tensor]]]]]:
         """
         Disentanglement cannot be processed on different subjects at the
         same time, we need to load `pr` subjects one by one. The batch