Trying to access global rank or world size freezes my script

[EvanZ]

I've been trying to figure out how to get IterableDataset to work in a multi-GPU instance, but I'm very confused by the discussions and code examples I've seen, so that has been difficult. But even something very basic like just trying to print out the GPU rank during training seems to give me problems:

    def training_step(self, batch, batch_idx):
        x = batch['src']
        y = batch['label'].double()
        mask = batch['mask']
        print(self.trainer.global_rank)

I have training_step defined in my LightningModel class. Just adding this one print statement freezes the script as soon as it looks like it's going to initialize the processes. You can see where I hit CMD+C here:

Size of training set: 27019
Size of validation set: 4766
==========================================================================================
Layer (type:depth-idx)                                            Param #
==========================================================================================
LitModel                                                          --
├─Embedding: 1-1                                                  998,000
├─TransformerEncoder: 1-2                                         --
│    └─ModuleList: 2-1                                            --
│    │    └─TransformerEncoderLayer: 3-1                          143,812
│    │    └─TransformerEncoderLayer: 3-2                          143,812
│    │    └─TransformerEncoderLayer: 3-3                          143,812
│    │    └─TransformerEncoderLayer: 3-4                          143,812
│    │    └─TransformerEncoderLayer: 3-5                          143,812
│    │    └─TransformerEncoderLayer: 3-6                          143,812
├─Linear: 1-3                                                     1,007,980
==========================================================================================
Total params: 2,868,852
Trainable params: 2,868,852
Non-trainable params: 0
==========================================================================================
GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
HPU available: False, using: 0 HPUs
2023/04/06 20:43:18 WARNING mlflow.utils.autologging_utils: You are using an unsupported version of pytorch. If you encounter errors during autologging, try upgrading / downgrading pytorch to a supported version, or try upgrading MLflow.
/home/ubuntu/deep-behavior-embedding/.virtualenv/lib/python3.7/site-packages/pytorch_lightning/loops/utilities.py:97: PossibleUserWarning: `max_epochs` was not set. Setting it to 1000 epochs. To train without an epoch limit, set `max_epochs=-1`.
  category=PossibleUserWarning,
[rank: 0] Global seed set to 123
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/4
^CTraceback (most recent call last):
Traceback (most recent call last):
  File "/home/ubuntu/.vscode-server/extensions/ms-python.python-2023.4.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/pydevd.py", line 3489, in <module>
  File "/home/ubuntu/.vscode-server/extensions/ms-python.python-2023.4.1/pythonFiles/lib/python/debugpy/_vendored/pydevd/pydevd.py", line 3489, in <module>
Traceback (most recent call last):

Am I calling it wrong or calling it from an incorrect location that would cause this behavior?