Update modeling_rwkv.py

modeling_rwkv.py

@@ -1043,6 +1043,7 @@ class RWKV(MyModule):
             elif embs is not None and tokens is not None:
                 seq_mode = len(tokens) > 1
                 x = w['emb.weight'][tokens if seq_mode else tokens[0]]
+                x = x.to(device=embs.device, dtype=embs.dtype, non_blocking=True)
                 x = torch.cat([x, embs], dim=0)
             else:
                 raise ValueError('Either tokens or embs must be provided')