Thank you for your prompt response @ptrblck ! I was actually trying to understand if there is any bottleneck in the code. This one is fairly simple example, I printed out execution time of each operation in each batch. It seems It takes ridiculously long time in computation of gradients(~10sec). The very same operation takes a fraction of a second if no batch is used.
I was wondering if you can provide a further insight, thank you!