From c231239fdfb6e3c484f091c6ea004f767ee2ea1f Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 30 Aug 2024 03:47:52 +0200
Subject: [PATCH] cleanup

Perf difference from `torch.tensor(tensor.data)` doesn't seem realistically measurable and the warning is annoying.
---
 dequant.py | 11 +++++------
 ops.py     | 13 +++++--------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/dequant.py b/dequant.py
index 454bd8c..3688d7f 100644
--- a/dequant.py
+++ b/dequant.py
@@ -4,20 +4,19 @@ import torch
 from tqdm import tqdm
 
 def dequantize_tensor(tensor, dtype=None, dequant_dtype=None):
-    data = torch.tensor(tensor.data)
     qtype = getattr(tensor, "tensor_type", None)
-    oshape = getattr(tensor, "tensor_shape", data.shape)
+    oshape = getattr(tensor, "tensor_shape", tensor.shape)
 
     if qtype in (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16):
-        return data.to(dtype)
+        return tensor.to(dtype)
     elif qtype in dequantize_functions:
         dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype
-        return dequantize(data, qtype, oshape, dtype=dequant_dtype).to(dtype)
+        return dequantize(tensor, qtype, oshape, dtype=dequant_dtype).to(dtype)
     else:
         # this is incredibly slow
         tqdm.write(f"Falling back to numpy dequant for qtype: {qtype}")
-        new = gguf.quants.dequantize(data.cpu().numpy(), qtype)
-        return torch.from_numpy(new).to(data.device, dtype=dtype)
+        new = gguf.quants.dequantize(tensor.cpu().numpy(), qtype)
+        return torch.from_numpy(new).to(tensor.device, dtype=dtype)
 
 def dequantize(data, qtype, oshape, dtype=None):
     """
diff --git a/ops.py b/ops.py
index dbdaed9..43f6299 100644
--- a/ops.py
+++ b/ops.py
@@ -71,20 +71,17 @@ class GGMLLayer(torch.nn.Module):
 
     def _save_to_state_dict(self, destination, prefix, keep_vars):
         # This is a fake state dict for vram estimation
-        if self.weight is not None:
-            weight = torch.zeros_like(self.weight, device=torch.device("meta"))
-            destination[f"{prefix}weight"] = weight
+        weight = torch.zeros_like(self.weight, device=torch.device("meta"))
+        destination[prefix + "weight"] = weight
         if self.bias is not None:
             bias = torch.zeros_like(self.bias, device=torch.device("meta"))
-            destination[f"{prefix}bias"] = bias
+            destination[prefix + "bias"] = bias
         return
 
         # This would return the actual state dict
-        weight, bias = self.get_weights()
-        if weight is not None:
-            destination[f"{prefix}weight"] = weight
+        destination[prefix + "weight"] = self.get_weight(self.weight)
         if bias is not None:
-            destination[f"{prefix}bias"] = weight
+            destination[prefix + "bias"] = self.get_weight(self.bias)
 
     def get_weight(self, tensor, dtype):
         if tensor is None: