A code snippet I could run on my machine to reproduce the large memory allocation.
At the moment I’m using your model definition and a random input in the shape [5, 1, 224, 224] in the forward pass, which only uses 120MB, which is far off of the reported 105GB.