fix tests to work with torch 1.8.0

## 🐛 Bug

A few tests do not pass when using torch v1.8.0

tests/nn/pipe_process/test_rpc.py::test_construct_only_rank_zero
tests/optim/test_oss.py::test_collect_shards
tests/optim/test_oss.py::test_reproducibility

## Command

## To Reproduce

Steps to reproduce the behavior:

1.
2.
3.

## Expected behavior

## Environment

Please copy and paste the output from the
environment collection script from PyTorch
(or fill out the checklist below manually).

You can run the script with:
```
# For security purposes, please check the contents of collect_env.py before running it.
python -m torch.utils.collect_env
```

- 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:

## Additional context