added code to turn value into tensor if not already

Xilinx · May 1, 2024 · 0428c34 · 0428c34
1 parent 165b8de
commit 0428c34
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/brevitas/nn/mixin/base.py b/src/brevitas/nn/mixin/base.py
@@ -79,7 +79,9 @@ 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))
- elif isinstance(inp, Tensor):
+ else:
+ if isinstance(inp, Tensor) is False:
+ inp = torch.tensor(inp)
  inp = inp.rename(None)
  return inp