DeepSpeed

5e3a0707 - pass only tensors to z3 hooks to prevent dynamo from displaying errors