netloader 3.11.3__tar.gz → 3.11.4__tar.gz

{netloader-3.11.3 → netloader-3.11.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: netloader
-Version: 3.11.3
+Version: 3.11.4
 Summary: Utility to generate and train PyTorch neural network objects from JSON files
 Author-email: Ethan Tregidga <ethan.tregidga@epfl.ch>
 License-Expression: MIT

{netloader-3.11.3 → netloader-3.11.4}/netloader/__init__.py

@@ -6,7 +6,7 @@ import logging
 import warnings
 
 
-__version__ = '3.11.3'
+__version__ = '3.11.4'
 __author__ = 'Ethan Tregidga'
 logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.WARNING)
 warnings.filterwarnings('once', category=DeprecationWarning, module=r'^netloader(\.|$)')

{netloader-3.11.3 → netloader-3.11.4}/netloader/architectures/flows.py

@@ -251,7 +251,7 @@ class NormFlowEncoder(BaseEncoder):
             'max': transform,
             'meds': transform,
         }
-        assert isinstance(self.net, (nn.ModuleList, Network))
+        assert isinstance(self.net, Network)
 
         if not self._train_encoder:
             self.net.layers[:-1].requires_grad_(False)

{netloader-3.11.3 → netloader-3.11.4}/netloader/models/misc.py

@@ -12,7 +12,7 @@ from netloader.utils import Config, TypedModuleList
 from netloader.network import BaseNetwork, CompatibleNetwork, Network
 
 
-class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
+class MultiNetwork(BaseNetwork[TypedModuleList[BaseNetwork]]):
     """
     Network class to hold multiple networks and train them together.
 
@@ -27,7 +27,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
         is the batch size
     kl_loss : Tensor
         KL divergence loss on the latent space of shape (1) and type float, if using a sample layer
-    layers : TypedModuleList[CompatibleNetwork | Network]
+    layers : TypedModuleList[BaseNetwork]
         Network construction
     """
     _cache: bool
@@ -36,14 +36,14 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
     def __init__(
             self,
             name: str,
-            *nets: nn.Module | nn.ModuleList | CompatibleNetwork | Network,
+            *nets: nn.Module | nn.ModuleList | BaseNetwork,
             save_outputs: bool = ...) -> None: ...
 
     @overload
     def __init__(
             self,
             name: str,
-            *nets: str | nn.Module | nn.ModuleList | CompatibleNetwork | Network | Config,
+            *nets: str | nn.Module | nn.ModuleList | BaseNetwork | Config,
             save_outputs: bool = ...,
             suppress_warning: bool = ...,
             root: str = ...,
@@ -55,7 +55,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
     def __init__(
             self,
             name: str,
-            *nets: str | nn.Module | nn.ModuleList | CompatibleNetwork | Network | Config,
+            *nets: str | nn.Module | nn.ModuleList | BaseNetwork | Config,
             save_outputs: bool = False,
             suppress_warning: bool = False,
             root: str = '',
@@ -68,7 +68,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
         ----------
         name : str
             Name of the network, used for saving
-        *nets: str | nn.Module | nn.ModuleList | CompatibleNetwork | Network | Config
+        *nets: str | nn.Module | nn.ModuleList | BaseNetwork | Config
             Networks to combine sequentially
         save_outputs : bool, Optional
             If outputs from each network should be saved, default = False
@@ -91,7 +91,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
         self.name = name
         self.layers = TypedModuleList()
         calc_in_shape: bool = True
-        net: str | nn.Module | nn.ModuleList | CompatibleNetwork | Network | Config
+        net: str | nn.Module | nn.ModuleList | BaseNetwork | Config
         builds: list[bool] = [isinstance(net, (str, Config)) for net in nets]
         out_shape: list[int] | None
         default: dict[str, Any] | None
@@ -150,15 +150,15 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
                 in_shape = out_shape
 
     @overload
-    def __getitem__(self, item: int) -> CompatibleNetwork | Network: ...
+    def __getitem__(self, item: int) -> BaseNetwork: ...
 
     @overload
-    def __getitem__(self, item: slice) -> TypedModuleList[CompatibleNetwork | Network]: ...
+    def __getitem__(self, item: slice) -> TypedModuleList[BaseNetwork]: ...
 
     def __getitem__(
             self,
             item: int | slice,
-    ) -> CompatibleNetwork | Network | TypedModuleList[CompatibleNetwork | Network]:
+    ) -> BaseNetwork | TypedModuleList[BaseNetwork]:
         """
         Returns the network(s) at the given index or slice.
 
@@ -169,7 +169,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
 
         Returns
         -------
-        CompatibleNetwork | Network | TypedModuleList[CompatibleNetwork | Network]
+        BaseNetwork | TypedModuleList[BaseNetwork]
             The network(s) at the given index or slice
         """
         return self.layers[item]
@@ -200,7 +200,7 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
         TensorListLike
             Output tensor from the network with shape (N,...) and type float
         """
-        net: CompatibleNetwork | Network
+        net: BaseNetwork
         self.checkpoints = []
 
         for i, net in enumerate(self.layers):
@@ -231,11 +231,14 @@ class MultiNetwork(BaseNetwork[TypedModuleList[CompatibleNetwork | Network]]):
         dict[str, Any] | Config
             Network configuration
         """
-        net: CompatibleNetwork | Network
+        net: BaseNetwork
         config = Config(layers=[net.get_config() for net in self.layers])
         return config.to_dict() if dict_ else config
 
     def to(self, *args: Any, **kwargs: Any) -> Self:
         super().to(*args, **kwargs)
-        self.layers.to(*args, **kwargs)
+        layer: BaseNetwork
+
+        for layer in self.layers:
+            layer.to(*args, **kwargs)
         return self

{netloader-3.11.3 → netloader-3.11.4}/netloader/network.py

@@ -203,7 +203,14 @@ class BaseNetwork(nn.Module, ABC, Generic[ModuleListT]):
 
     def to(self, *args: Any, **kwargs: Any) -> Self:
         super().to(*args, **kwargs)
-        self.kl_loss = self.kl_loss.to(*args, **kwargs)
+        self.kl_loss = self.kl_loss.detach().to(*args, **kwargs)
+
+        for i, checkpoint in enumerate(self.checkpoints):
+            if hasattr(checkpoint, 'to'):
+                self.checkpoints[i] = checkpoint.detach().to(*args, **kwargs)
+            else:
+                self._logger.warning(f'Failed to move checkpoint {checkpoint.__class__.__name__} '
+                                     f'to device')
         return self
 
 
@@ -544,13 +551,6 @@ class Network(BaseNetwork[TypedModuleList[layers.BaseLayer]]):
 
         for layer in self.layers:
             layer.to(*args, **kwargs)
-
-        for i, checkpoint in enumerate(self.checkpoints):
-            if hasattr(checkpoint, 'to'):
-                self.checkpoints[i] = checkpoint.to(*args, **kwargs)
-            else:
-                self._logger.warning(f'Failed to move checkpoint {checkpoint.__class__.__name__} '
-                                     f'to device')
         return self
 
 

{netloader-3.11.3 → netloader-3.11.4}/netloader.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: netloader
-Version: 3.11.3
+Version: 3.11.4
 Summary: Utility to generate and train PyTorch neural network objects from JSON files
 Author-email: Ethan Tregidga <ethan.tregidga@epfl.ch>
 License-Expression: MIT