Fix Imports

Xilinx · Jan 22, 2024 · e82364f · e82364f
1 parent 058ba0c
commit e82364f
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/brevitas_ort/quant_module_cases.py b/tests/brevitas_ort/quant_module_cases.py
@@ -5,6 +5,9 @@
 from pytest_cases import set_case_id
 from torch import nn
 
+from brevitas.nn.quant_activation import QuantIdentity
+from brevitas.nn.quant_avg_pool import TruncAvgPool2d
+from brevitas.nn.quant_rnn import QuantLSTM
 from brevitas.quant.scaled_int import Int32Bias
 
 from .common import *