pyproject.toml

@@ -3,7 +3,7 @@ requires = ["setuptools>=45", "wheel"]
 build-backend = "setuptools.build_meta"
 [project]
 name = "aiunn"
-version = "0.4.0"
+version = "0.4.1"
 description = "Finetuner for image upscaling using AIIA"
 readme = "README.md"
 requires-python = ">=3.10"

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
     name="aiunn",
-    version="0.4.0",
+    version="0.4.1",
     packages=find_packages(where="src"),
     package_dir={"": "src"},
     install_requires=[

src/aiunn/__init__.py

@@ -4,4 +4,4 @@ from .upsampler.aiunn import aiuNN
 from .upsampler.config import aiuNNConfig
 from .inference.inference import aiuNNInference
 
-__version__ = "0.4.0"
+__version__ = "0.4.1"

src/aiunn/finetune/memory_trainer.py

@@ -177,6 +177,7 @@ class MemoryOptimizedTrainer(aiuNNTrainer):
                 
                 with autocast(device_type=self.device.type):
                     outputs = self.model(low_res)
+                    outputs = outputs.clone()
                     loss = self.criterion(outputs, high_res)
                 
                 val_loss += loss.item()
@@ -252,10 +253,10 @@ class MemoryOptimizedTrainer(aiuNNTrainer):
                     if hasattr(self, 'use_checkpointing') and self.use_checkpointing:
                         low_res.requires_grad_()
                         outputs = checkpoint(self.model, low_res)
-                        outputs = outputs.clone()  # <-- Clone added here
+                        outputs = outputs.clone()
                     else:
                         outputs = self.model(low_res)
-                        outputs = outputs.clone()  # <-- Clone added here
+                        outputs = outputs.clone()
                     loss = self.criterion(outputs, high_res)
                     
                     # Scale loss for gradient accumulation