I am trying to create a layer that is executed every batch in a layer using the torch.
Previously, I tried to write it with TensorFlow and Keras, but NumPy was not possible, so I used PyTorch. When I used NumPy in pytroch, it was run with cpu and it was too slow.
So, in the end, I wrote NumPy by replacing it with the function of the torch.
I am currently training, but it is too slow and the accuracy is 10%. (take 5hour per a epoch)
The sequence of the code is as follows:
class my_custom_layer(nn.Module):
    def __init__(self, D_in, D_out, kernel_size):
        super(my_custom_layer, self).__init__()
        self.BConv_layer = nn.Conv2d(D_in, D_out, kernel_size, padding=(1))
    def forward(self, x):      
        with torch.no_grad(): 
            y = my_function(x)
            #preprocess batch using histogram data
            #make histogram of img -> Partially scaling according to the characteristics of the image -> return scaling factor for result of conv layer

        results = F.conv(x, kernel)
        results = torch.multiply(results, y)
        return results

Sometimes it runs on the CPU, but most of the time it runs on the GPU.
I will replace this layer with VGG’s conv layer.
Is it correct to write a custom layer like this?