Merging Convolutional & Batch Norm Layer

Hi. I’m trying to merge a convolutional layer and a batch norm layer so I dont have to create a separate representation of the batch norm layer in an FPGA implementation of the model. I merged the layers by replacing the weights and biases of the preceding convolutional layer with the appropriate transformed weights and biases required for the fusion.

The problem is that I cannot perform in-place operations on leaf-type variables i.e. the weights and biases, so I have to use torch.no_grad() while testing the model. However, that way, I can’t train it and the loss is never changing because the model can’t backpropagate. Any help on this would be appreciated.