DeepSpeed

b3870363 - [Stage][Fix] Add additional conditions when checking types of output from the model (#1026)

[Stage][Fix] Add additional conditions when checking types of output from the model (#1026) * Add additional conditions when checking types of output from the model * Add test * Modify test to use torch.tensor as well Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>