Torch profiler: memory consumption on particular GPU

Hello!

Currently seems that torch.autograd.profile memory profiler prints total amount of CUDA tensors on ALL used GPUs. I want to get a dictionary {device_id : memory}. How can i do this?

Thank you!