Error in distributed test

`test_distributed.test_squared_edges_are_same_dist_split_over_dataset` is failing with the following error:
```
ValueError: Invalid input: /tmp/pytest-of-user/pytest-6/test_squared_edges_are_same_di0/double_out/distributed_test_double is neither a .pt file nor a directory
```