Update lit_llama/utils.py

lit_llama/utils.py

@@ -89,13 +89,13 @@ class EmptyInitOnDevice(torch.overrides.TorchFunctionMode):
         if self.quantization_mode == 'llm.int8':
             if device.type != "cuda":
                 raise ValueError("Quantization is only supported on the GPU.")
-            from lit_llama.quantization import Linear8bitLt
+            from .quantization import Linear8bitLt
             self.quantized_linear_cls = Linear8bitLt
         elif self.quantization_mode == 'gptq.int4':
-            from lit_llama.quantization import ColBlockQuantizedLinear
+            from .quantization import ColBlockQuantizedLinear
             self.quantized_linear_cls = functools.partial(ColBlockQuantizedLinear, bits=4, tile_cols=-1)
         elif self.quantization_mode == 'gptq.int8':
-            from lit_llama.quantization import ColBlockQuantizedLinear
+            from .quantization import ColBlockQuantizedLinear
             self.quantized_linear_cls = functools.partial(ColBlockQuantizedLinear, bits=8, tile_cols=-1)
         elif self.quantization_mode is not None:
             raise RuntimeError(f"unknown quantization mode {self.quantization_mode}")