Add support for multi node for Optimize & StreamingDataset

[tchaton]

🐛 Bug

To Reproduce

Steps to reproduce the behavior...

Code sample

# Ideally attach a minimal code sample to reproduce the decried issue.
# Minimal means having the shortest code but still preserving the bug.

Expected behavior

Additional context

Environment detail

• PyTorch Version (e.g., 1.0):
• OS (e.g., Linux):
• How you installed PyTorch (conda, pip, source):
• Build command you used (if compiling from source):
• Python version:
• CUDA/cuDNN version:
• GPU models and configuration:
• Any other relevant information: