added isinstance elif incase variable isn't a tensor

Xilinx · May 1, 2024 · 165b8de · 165b8de
1 parent 28e8868
commit 165b8de
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/brevitas/nn/mixin/base.py b/src/brevitas/nn/mixin/base.py
@@ -79,7 +79,7 @@ def unpack_input(self, inp: Union[Tensor, QuantTensor]) -> Union[Tensor, QuantTe
         if not torch._C._get_tracing_state():
             if isinstance(inp, QuantTensor):
                 inp = inp.set(value=inp.value.rename(None))
-            else:
+            elif isinstance(inp, Tensor):
                 inp = inp.rename(None)
         return inp