diff options
Diffstat (limited to 'models/model.py')
| -rw-r--r-- | models/model.py | 10 |
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 |