From 505b3836054f3ab45d0784d7f613c92b770d533e Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Thu, 3 Jul 2025 14:22:08 +0200
Subject: [PATCH] corrected MOTrainer

---
 pyproject.toml                       | 2 +-
 setup.py                             | 2 +-
 src/aiunn/__init__.py                | 2 +-
 src/aiunn/finetune/memory_trainer.py | 5 +++--
 4 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 64b4591..f021b09 100644
--- a/pyproject.toml
+++ b/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"
diff --git a/setup.py b/setup.py
index f950e60..736b66f 100644
--- a/setup.py
+++ b/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=[
diff --git a/src/aiunn/__init__.py b/src/aiunn/__init__.py
index 2b68267..59d575f 100644
--- a/src/aiunn/__init__.py
+++ b/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"
\ No newline at end of file
+__version__ = "0.4.1"
\ No newline at end of file
diff --git a/src/aiunn/finetune/memory_trainer.py b/src/aiunn/finetune/memory_trainer.py
index 9a7bd8f..26a1b8d 100644
--- a/src/aiunn/finetune/memory_trainer.py
+++ b/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