If the input doesn’t require gradients, calculating the gradients for a single layer should be faster than all layers.
You code doesn’t show any profiling so I cannot comment on the correctness of it.
Note that, in case you are using a GPU, CUDA operations are asynchronous, so you would have to synchronize the code before starting and stopping the timer or use the torch.utils.benchmark utility from the nightly binaries.
That being said, you are currently also synchronizing your “production” code by pushing the gradient norm to the CPU and calling item() on it.