but the big absolute overhead you’ve observed is from repeated Tensor unwrappings in the loop. If your neural network iteration consists of 100-1000 operations, and tensor sizes are bigger (from batching), this overhead is not so significant.
for long sequences of piecewise operations, JIT compiling does a decent job fusing them (it used to need USE_CPU_FUSER=1 env. variable, haven’t tried it recently)