ptrblck, thanks for your inputs.
Here again the debug info before and after calling clip_feature function running on multiple GPUs.
EncoderDecoder clip_feature masks.shape in (38, 61)
masks.device:cuda:0
EncoderDecoder clip_feature masks.shape in (38, 61)
masks.device:cuda:1
EncoderDecoder clip_feature masks.shape in (38, 61)
masks.device:cuda:2
max_len:50      max_len.device: cuda:0
EncoderDecoder clip_feature masks.shape in (36, 61)
max_len:54       max_len.device: cuda:1
max_len:61      max_len.device: cuda:2
masks.device:cuda:3
max_len:51      max_len.device: cuda:3

It’s actually a big jump to show where the error comes since there are lots of operations before calling the relation_geo_attention function. Since the whole snippet works on single GPU, I guess it’s necessary to poll all the codes out here (hopefully I’m right:D). The trigger is the dimension after clip_feature is totally wrong away from the expectation (as one output by single GPU).  The error comes from lines marked with ^^^^s.
def relation_geo_attention(query, key, value, box_embd_matrix, mask=None):
    N = value.size()[:2]
    dim_k = key.size(-1)
    dim_g = box_embd_matrix.size()[-1]

    w_q = query
    w_k = key.transpose(-2, -1)
    w_v = value
    w_g = box_embd_matrix
	
    #attention weights
    scaled_dot = torch.matmul(w_q,w_k)
    w_a = scaled_dot / np.sqrt(dim_k)
    if mask is not None:
        w_a = w_a.masked_fill(mask == 0, -1e9)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        #! RuntimeError occurs from here

    # calculating  retlation between geometric and feature
    w_mn = torch.log(torch.clamp(w_g, min = 1e-6)) + w_a
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    #! RuntimeError occurs possibly here either

After clipping features, the returned dimension is not as expected, all the following operations in transformer (the relation geometric attention function) are wrong. I do not think it’s a scatter problem from devices.
RuntimeError: Caught RuntimeError in replica 0 on device 0.
Original Traceback (most recent call last):
....
....
w_a = w_a.masked_fill(mask == 0, -1e9)
RuntimeError: The size of tensor a (61) must match the size of tensor b (50) at non-singleton dimension 3

What really confuses me is why it works on single device, but got errors instead of running it on n-GPUs.
Thanks
Anakin