homa 0.2.93__tar.gz → 0.2.96__tar.gz

{homa-0.2.93 → homa-0.2.96}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: homa
-Version: 0.2.93
+Version: 0.2.96
 Summary: A curated list of machine learning and deep learning helpers.
 Author-email: Taha Shieenavaz <tahashieenavaz@gmail.com>
 Requires-Python: >=3.7

{homa-0.2.93 → homa-0.2.96}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "homa"
-version = "0.2.93"
+version = "0.2.96"
 description = "A curated list of machine learning and deep learning helpers."
 authors = [
   { name="Taha Shieenavaz", email="tahashieenavaz@gmail.com" },

{homa-0.2.93 → homa-0.2.96}/src/homa/ensemble/concerns/StoresModels.py

@@ -1,7 +1,8 @@
 import torch
-from copy import deepcopy
+import io
 from typing import List
 from ...vision import Model
+from ..utils import get_model_device
 
 
 class StoresModels:
@@ -12,11 +13,17 @@ class StoresModels:
     def record(self, model: Model | torch.nn.Module):
         model_: torch.nn.Module | None = None
         if isinstance(model, Model):
-            model_ = deepcopy(model.network)
+            model_ = model.network
         elif isinstance(model, torch.nn.Module):
-            model_ = deepcopy(model)
+            model_ = model
         else:
             raise TypeError("Wrong input to ensemble record")
+
+        device = get_model_device(model_)
+        buffer = io.BytesIO()
+        torch.save(model_.to("cpu"), buffer)
+        buffer.seek(0)
+        model_ = torch.load(buffer, map_location=device)
         self.models.append(model_)
 
     def push(self, *args, **kwargs):

homa-0.2.96/src/homa/ensemble/utils.py

@@ -0,0 +1,9 @@
+import torch
+
+
+def get_model_device(model: torch.nn.Module):
+    try:
+        device = next(model.parameters()).device
+    except StopIteration:
+        device = torch.device("cpu")
+    return device

{homa-0.2.93 → homa-0.2.96}/src/homa.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: homa
-Version: 0.2.93
+Version: 0.2.96
 Summary: A curated list of machine learning and deep learning helpers.
 Author-email: Taha Shieenavaz <tahashieenavaz@gmail.com>
 Requires-Python: >=3.7

{homa-0.2.93 → homa-0.2.96}/src/homa.egg-info/SOURCES.txt

@@ -81,6 +81,7 @@ src/homa/cli/namespaces/MakeNamespace.py
 src/homa/cli/namespaces/__init__.py
 src/homa/ensemble/Ensemble.py
 src/homa/ensemble/__init__.py
+src/homa/ensemble/utils.py
 src/homa/ensemble/concerns/CalculatesMetricNecessities.py
 src/homa/ensemble/concerns/PredictsProbabilities.py
 src/homa/ensemble/concerns/ReportsClassificationMetrics.py