Model on FPGA Outputting Unexpected Result

[Phu1234567]

I am training my Brevitas model on a radio modulation dataset with input shape (1, 2, 1024) and the raw dataset was converted into uint8 (0-255). The model reached 84% accuracy right after training. But after running the building step, and the driver is moved onto the FPGA, it gives me unusual outputs. I am suspecting the problem has to do with the datatype of my dataset but I do not know how to confirm it. The dataset I trained on the Brevitas model is the same one I transported onto the FPGA for validation. The only difference between them is their shape having an extra axis after reshaping to fit with the driver requirement.

My driver.py:

My custom validate.py script on the FPGA

The output

How I define my model (taken from the radio sandbox from FINN example, with minor modification for my dataset)

from torch import nn
import brevitas.nn as qnn
from brevitas.quant import Int16Bias
from brevitas.inject.enum import ScalingImplType
from brevitas.inject.defaults import Int8ActPerTensorFloatMinMaxInit

# Adjustable hyperparameters
input_bits = 8
a_bits = 8
w_bits = 8
filters_conv = 64
filters_dense = 128

# Setting seeds for reproducibility
torch.manual_seed(0)
np.random.seed(0)

class InputQuantizer(Int8ActPerTensorFloatMinMaxInit):
    bit_width = input_bits
    min_val = 0.0
    max_val = 255.0
    scaling_impl_type = ScalingImplType.CONST # Fix the quantization range to [min_val, max_val]

class QuantizedVGG(nn.Module):
    def __init__(self):
        super(QuantizedVGG, self).__init__()
        self.layers=nn.Sequential(
            # Input quantization layer
            qnn.QuantHardTanh(act_quant=InputQuantizer),

            qnn.QuantConv1d(2, filters_conv, 3, padding=1, weight_bit_width=w_bits, bias=False),
            nn.BatchNorm1d(filters_conv),
            qnn.QuantReLU(bit_width=a_bits),
            nn.MaxPool1d(2),

            qnn.QuantConv1d(filters_conv, filters_conv, 3, padding=1, weight_bit_width=w_bits, bias=False),
            nn.BatchNorm1d(filters_conv),
            qnn.QuantReLU(bit_width=a_bits),
            nn.MaxPool1d(2),

            qnn.QuantConv1d(filters_conv, filters_conv, 3, padding=1, weight_bit_width=w_bits, bias=False),
            nn.BatchNorm1d(filters_conv),
            qnn.QuantReLU(bit_width=a_bits),
            nn.MaxPool1d(2),

            qnn.QuantConv1d(filters_conv, filters_conv, 3, padding=1, weight_bit_width=w_bits,bias=False),
            nn.BatchNorm1d(filters_conv),
            qnn.QuantReLU(bit_width=a_bits),
            nn.MaxPool1d(2),

            qnn.QuantConv1d(filters_conv, filters_conv, 3, padding=1, weight_bit_width=w_bits, bias=False),
            nn.BatchNorm1d(filters_conv),
            qnn.QuantReLU(bit_width=a_bits),
            nn.MaxPool1d(2),

            qnn.QuantConv1d(filters_conv, filters_conv, 3, padding=1, weight_bit_width=w_bits, bias=False),
            nn.BatchNorm1d(filters_conv),
            qnn.QuantReLU(bit_width=a_bits),
            nn.MaxPool1d(2),

            qnn.QuantConv1d(filters_conv, filters_conv, 3, padding=1, weight_bit_width=w_bits, bias=False),
            nn.BatchNorm1d(filters_conv),
            qnn.QuantReLU(bit_width=a_bits),
            nn.MaxPool1d(2),
            
            nn.Flatten(),

            qnn.QuantLinear(filters_conv*8, filters_dense, weight_bit_width=w_bits, bias=False),
            nn.BatchNorm1d(filters_dense),
            qnn.QuantReLU(bit_width=a_bits, return_quant_tensor=True),

            qnn.QuantLinear(filters_dense, filters_dense, weight_bit_width=w_bits, bias=False),
            nn.BatchNorm1d(filters_dense),
            qnn.QuantReLU(bit_width=a_bits, return_quant_tensor=True),

            qnn.QuantLinear(filters_dense, 27, weight_bit_width=w_bits, bias=True, bias_quant=Int16Bias),
        )

    def forward(self, x):
        x=self.layers(x)
        return x
    
model=QuantizedVGG()
model=model.to('cuda')

My build step (I skipped the step_set_fifo_depths):

# assemble build flow from custom and pre-existing steps
def select_build_steps(platform):
    return [
        "step_tidy_up",
        step_pre_streamline,
        "step_streamline",
        "step_convert_to_hw",
        step_convert_final_layers,
        "step_create_dataflow_partition",
        "step_specialize_layers",
        "step_target_fps_parallelization",
        "step_apply_folding_config",
        "step_minimize_bit_width",
        "step_generate_estimate_reports",
        "step_hw_codegen",
        "step_hw_ipgen",
        # "step_set_fifo_depths",
        "step_create_stitched_ip",
        "step_measure_rtlsim_performance",
        "step_out_of_context_synthesis",
        "step_synthesize_bitfile",
        "step_make_pynq_driver",
        "step_deployment_package",
    ]