@mcarilli I can confirm this also happens with 1.6.0.dev20200430. Also it affects LSTM as well (my case), while the manually permuting inputs and changing batch_first doesn’t help (in the past some things were broken with batch_first=True).
Also, not sure how to reach you regarding specific use case with variable batch size. So I read materials on AMP and to squeeze out max performance all dimensions on layers like Linear or Conv2D must be divisible by 8 and batch_size as well. Is there any kind of workaround to get decent performance with variable batch size? I can ask completely new question if that is preferred.