[TiledMLP]: fix for bs>1 (#7412)
It looks like my TiledMLP was working correctly only for batch_size=1
fixing to work with any bs
thanks to @winglian for detecting the problem and sending me an easy
repro
---------
Signed-off-by: Stas Bekman <stas@stason.org>