Solved by albanD in post #6 
              
              
                Fun fact, the lambda is actually significantly faster than creating a function and reusing it  
In [1]: import torch

In [2]: def gs(t):
   ...:     return t.size()
   ...:

In [3]: t = torch.rand(10, 3, 2)

In [4]: %timeit gs(t)
692 ns ± 12.6 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops e…