Thanks for your reply again, ptrblck.
I think probably you got confused by my description.  As so far, the data is correctly split in the chunks
[38 x 61] [38 x 61] [38 x 61] [36 x 61]

The input length arises the error. Concretely, I have a relation model that calculates relation between geometric and appearance. This relation model in organized  in a ModuleList in cascade fashion. The attentions are fed to the first module in ModuleList. Then its output will be fed to following modules further.
If the data just scattered on single device, all the relation models have the aligned (same) dimension. It works. If the data got chunked and scattered on 4 devices, the features have different lengths on single device, i.e, device 0 as debug info suggests below.
     if mask is not None:
            print('mask.size', mask.size())
            print('mask.device:\t%s' % mask.device) 
            print('w_a.size', w_a.size())
            print('w_a.device:\t%s' % w_a.device)

            assert query.device == key.device, 'query and key are not on the same device'
            assert value.device == key.device, 'value and key are not on the same device'
            assert query.device == box_relation_embds_matrix.device, 'query and box are not on the same device'
            assert query.device == mask.device, 'query and mask are not on the same device'
            assert w_a.device == mask.device, 'w_a and mask are not on the same device'

            w_a = w_a.masked_fill(mask == 0, -1e9)

All the assertions hold true.  Here the corresponding debug outputs(just for simplicity and convenience, i used two devices):
###clip_feature
EncoderDecoder clip_feature masks.shape in torch.Size([75, 61])
masks.device:cuda:0
EncoderDecoder clip_feature masks.shape in torch.Size([75, 61])
masks.device:cuda:1
max_len:54      max_len.device: cuda:0
max_len:61      max_len.device: cuda:1

### got padded somewhere later
padded.size torch.Size([75, 61, 512])
padded.size torch.Size([75, 54, 512])

### the first relation model info
 mask.size torch.Size([75, 1, 1, 61])
 mask.device:       cuda:1
 w_a.size torch.Size([75, 8, 61, 61])
 w_a.device:        cuda:1

### the second relation model info
 mask.size torch.Size([75, 1, 1, 61])
                       ^^^^^^^^^^^^^^
 mask.device:       cuda:0
 w_a.size torch.Size([75, 8, 54, 54])
                       ^^^^^^^^^^^^^^
 w_a.device:        cuda:0

As in last reply shows, the problem comes from the device 0, showing as follow:



 Anakin:

RuntimeError: Caught RuntimeError in replica 0 on device 0. Original Traceback (most recent call last):
error comes at this line
w_a = w_a.masked_fill(mask == 0, -1e9)


For comparison purpose, I post the log of single device:
# clip_feature
EncoderDecoder clip_feature masks.shape in torch.Size([150, 61])
masks.device:cuda:0
max_len:61      max_len.device: cuda:0
padded.size torch.Size([150, 61, 512])

### the first relation model info
mask.size torch.Size([150, 1, 1, 61])
mask.device:       cuda:0
w_a.size torch.Size([150, 8, 61, 61])
w_a.device:        cuda:0

### the same for other relation models

It seems like a coding problem rather than a scattering problem, as I guess as so far. I’ll dig it out continuously. Any inputs will be thankful.