Make data contiguous before the inplace reshape-copy_ function (#2489)

Co-authored-by: Michael Wyatt <[email protected]>

deepspeed/module_inject/load_checkpoint.py

@@ -15,6 +15,7 @@ def load_model_with_checkpoint(r_module,
     error_msgs = []
 
     def transpose(data):
+        data = data.contiguous()
         data1 = data.transpose(-1, -2).reshape(-1)
         data.reshape(-1).copy_(data1)
         data1 = None

deepspeed/module_inject/replace_module.py

@@ -215,6 +215,7 @@ def replace_attn(child, policy):
         attn_module = transformer_inference.DeepSpeedDiffusersAttention(config)
 
         def transpose(data):
+            data = data.contiguous()
             data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1))
             data = data.reshape(data.shape[-1], data.shape[-2])
             data.to(torch.cuda.current_device())
@@ -531,7 +532,7 @@ def replace_with_policy(child,
             # transpose it here to reduce inference cost!
             def transpose(data):
                 # temp move to cpu to avoid requiring extra GPU memory during the reshape
-                data = data.to('cpu')
+                data = data.to('cpu').contiguous()
                 data.reshape(-1).copy_(data.transpose(-1, -2).contiguous().reshape(-1))
                 data = data.reshape(data.shape[-1], data.shape[-2])
                 data.to(torch.cuda.current_device())