summaryrefslogtreecommitdiff
path: root/models/model.py
diff options
context:
space:
mode:
Diffstat (limited to 'models/model.py')
-rw-r--r--models/model.py10
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