Moving tensors to MPS device taking extremely long time

mps
1

Hi,

We are trying to run our code on mac and MPS, and the following code shows that just the line with .to(device) takes up to 50 seconds to move a tensor of shape (8,16,16,128,128).

with torch.no_grad():

        for num, data in enumerate(data_loader):

      start_loop = time.time()

      print(f'Processing: {num} of {dataset_test.__len__()}')

      start= time.time()

      input_data = data['image'].float().to(device)

      print("data to device time", time.time()-start)

      name_data = np.array(data['names'])

      input_data = torch.permute(input_data, (0, 2, 1, 3, 4))

      name_data = np.transpose(name_data)

      start = time.time()   

      pred_comb = model(input_data)

      print('Model time',time.time() - start)

On CPU of course the model is the bottleneck. But switching to MPS slows every single .to(device) or .cpu() calls. This is with latest pytorch version and on M1 chip. The dataset is a just a loading a npy file and returning a sequence of them. And dataloader:

data_loader = torch.utils.data.DataLoader( dataset_test,batch_size=batch_size,shuffle=False,num_workers=num_workers,pin_memory=False)

This is the print output :

Processing: 0 of 526
data to device time torch.Size([4, 16, 1, 128, 128]) 0.01835012435913086
Model time 0.5450329780578613
Loop time 0.5635440349578857
Processing: 1 of 526
data to device time torch.Size([4, 16, 1, 128, 128]) 7.373780965805054
Model time 0.004361867904663086
Loop time 7.378338098526001
Processing: 2 of 526
data to device time torch.Size([4, 16, 1, 128, 128]) 7.838125228881836
Model time 0.002288341522216797
Loop time 7.84054708480835
Processing: 3 of 526

So clearly the .to(device) is way to slow than it should.