If the tensor creation is slow in your code, you could calculate the weights once, store them on the GPU and index them using the current label tensor.
Your general approach look alright, but you should note that a weighted loss is usually normalized to avoid a dependency on the currently used class label distribution in the batch as described e.g. here.