replace types dict to Mapping

Author: bdvllrs

shimmer/modules/dict_buffer.py

@@ -1,4 +1,4 @@
-from collections.abc import Generator, Iterator
+from collections.abc import Generator, Iterator, Mapping
 from typing import Any
 
 import torch
@@ -7,7 +7,7 @@
 
 class DictBuffer(nn.Module):
     def __init__(
-        self, buffer_dict: dict[str, torch.Tensor], persistent: bool = True
+        self, buffer_dict: Mapping[str, torch.Tensor], persistent: bool = True
     ):
         super().__init__()
 

shimmer/modules/global_workspace.py

@@ -32,7 +32,7 @@ class GlobalWorkspaceBase(LightningModule):
     def __init__(
         self,
         gw_mod: GWModule,
-        domain_mods: dict[str, DomainModule],
+        domain_mods: Mapping[str, DomainModule],
         coef_buffers: DictBuffer,
         loss_mod: GWLosses,
         optim_lr: float = 1e-3,
@@ -271,7 +271,7 @@ def __init__(
         domain_mods: Mapping[str, DomainModule],
         gw_interfaces: Mapping[str, GWInterfaceBase],
         workspace_dim: int,
-        loss_coefs: dict[str, torch.Tensor],
+        loss_coefs: Mapping[str, torch.Tensor],
         optim_lr: float = 1e-3,
         optim_weight_decay: float = 0.0,
         scheduler_args: SchedulerArgs | None = None,
@@ -298,7 +298,7 @@ def __init__(
         domain_mods: Mapping[str, DomainModule],
         gw_interfaces: Mapping[str, GWInterfaceBase],
         workspace_dim: int,
-        loss_coefs: dict[str, torch.Tensor],
+        loss_coefs: Mapping[str, torch.Tensor],
         var_contrastive_loss: bool = False,
         optim_lr: float = 1e-3,
         optim_weight_decay: float = 0.0,
@@ -327,7 +327,7 @@ def pretrained_global_workspace(
     domain_mods: Mapping[str, DomainModule],
     gw_interfaces: Mapping[str, GWInterfaceBase],
     workspace_dim: int,
-    loss_coefs: dict[str, torch.Tensor],
+    loss_coefs: Mapping[str, torch.Tensor],
     var_contrastive_loss: bool = False,
     **kwargs,
 ) -> GlobalWorkspaceBase: