From 65a7c895bb0ac9547ba2f89d55fbdb609aa2bfe7 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sat, 14 Dec 2024 00:38:11 +0100
Subject: [PATCH] Better propagation fix

#163
---
 ops.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/ops.py b/ops.py
index db421cf..378d21e 100644
--- a/ops.py
+++ b/ops.py
@@ -119,7 +119,10 @@ class GGMLLayer(torch.nn.Module):
 
         # dequantize tensor while patches load
         weight = dequantize_tensor(tensor, dtype, self.dequant_dtype)
-        weight.__class__ = torch.Tensor # prevent propagating custom tensor class
+
+        # prevent propagating custom tensor class
+        if isinstance(weight, GGMLTensor):
+            weight.__class__ = torch.Tensor
 
         # apply patches
         if patch_list:
@@ -152,8 +155,14 @@ class GGMLLayer(torch.nn.Module):
 
     def forward_comfy_cast_weights(self, input, *args, **kwargs):
         if self.is_ggml_quantized():
-            return self.forward_ggml_cast_weights(input, *args, **kwargs)
-        return super().forward_comfy_cast_weights(input, *args, **kwargs)
+            out = self.forward_ggml_cast_weights(input, *args, **kwargs)
+        else:
+            out = super().forward_comfy_cast_weights(input, *args, **kwargs)
+
+        # non-ggml forward might still propagate custom tensor class
+        if isinstance(out, GGMLTensor):
+            out.__class__ = torch.Tensor
+        return out
 
     def forward_ggml_cast_weights(self, input):
         raise NotImplementedError