def get_current_gpu_memory_use():
            """returns a list of cuda memory allocations per GPU in MBs"""

            per_device_memory = []
            for id in range(torch.cuda.device_count()):
                with torch.cuda.device(id):
                    per_device_memory.append(torch.cuda.memory_allocated() >> 20)

            return per_device_memory
returns a list of cuda memory allocations per GPU in MBs
dpkd/transformers/tests/test_modeling_common.py