Solved by nullgeppetto in post #3 
              
              
                Hi @ptrblck, many thanks for your response. Based on your suggestion and the discussion found here, I figured it out. For the sake of completeness, I give a simple example below: 
import torch
import torch.nn as nn
import numpy as np

# Input tensor of shape (bs, C_in, H, W)
bs = 128
C_in = 3
H = 32…