RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

[thirttyyy]

Hi, the model is to large to out of my GPU meomery. So i want to use 2 GPUS to dispatch the model, and use the device_map like:

{'transformer.wte': 0, 'lm_head': 0, 'transformer.wpe': 0, 'transformer.drop': 0, 'transformer.h.0': 0, 'transformer.h.1': 0, 'transformer.h.2': 0, 'transformer.h.3': 0, 'transformer.h.4': 0, 'transformer.h.5': 0, 'transformer.h.6': 0, 'transformer.h.7': 0, 'transformer.h.8': 0, 'transformer.h.9': 0, 'transformer.h.10': 0, 'transformer.h.11': 0, 'transformer.h.12': 0, 'transformer.h.13': 0, 'transformer.h.14': 0, 'transformer.h.15': 0, 'transformer.h.16': 0, 'transformer.h.17': 0, 'transformer.h.18': 0, 'transformer.h.19': 0, 'transformer.h.20': 0, 'transformer.h.21': 0, 'transformer.h.22': 0, 'transformer.h.23': 0, 'transformer.h.24': 0, 'transformer.h.25': 0, 'transformer.h.26.ln_1': 0, 'transformer.h.26.attn': 0, 'transformer.h.26.ln_2': 0, 'transformer.h.26.mlp.c_fc': 0, 'transformer.h.26.mlp.c_proj': 1, 'transformer.h.26.mlp.act': 1, 'transformer.h.26.mlp.dropout': 1, 'transformer.h.27': 1, 'transformer.h.28': 1, 'transformer.h.29': 1, 'transformer.h.30': 1, 'transformer.h.31': 1, 'transformer.h.32': 1, 'transformer.h.33': 1, 'transformer.h.34': 1, 'transformer.h.35': 1, 'transformer.h.36': 1, 'transformer.h.37': 1, 'transformer.h.38': 1, 'transformer.h.39': 1, 'transformer.ln_f': 1}

It can load the model Successfully but got an error when inference:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

How can i fix it.