Move hooked embeddings to cpu to avoid oom (EleutherAI#34)

Co-authored-by: Samuel Weinbach <samuel.weinbach@gmail.com>

Author: sweinbach

deepspeed/runtime/engine.py

@@ -242,7 +242,7 @@ def hook_fn(module, input, output):
                     return
             else:
                 key = module.__class__.__name__
-            self.layer_outputs[key] = output 
+            self.layer_outputs[key] = [o.cpu() if torch.is_tensor(o) else o for o in output]
 
         def get_all_layers(net):
             for name, layer in net._modules.items():