Solved by albanD in post #2 
              
              
                Hi, 
register_buffer is used to save buffers that you then use during the forward pass. 
I would recommend using the autograd profiler to measure the runtime.