Both GPUs are getting same data from IterableDataset with DDP

[akshat-suwalka]

Hey Guys,
I am working on Large dataset which cannot be loaded in Memory. So I move to IterableDataset.
I have 1 node which have 2 gpus (T4) each of memory around 110GB
I am following the doc which tells how to IterableDataset for multi gpus. When i am applying for my code then i found that my both GPU are getting same data. Can anyone please help me to understand what i am missing?
Below is my code

`

class PowConDataSetIterable(IterableDataset):
    def __init__(self, s3_path, file_name, seq_len=14, batch_size=64):
        super().__init__()
        self.s3_path = s3_path
        self.start = 1
        self.end = 7
        self.seq_len = seq_len
        self.batch_size = batch_size

    def __iter__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:  # single-process data loading, return the full iterator
            iter_start = self.start
            iter_end = self.end
        else:  # in a worker process
            # split workload
            per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
            worker_id = worker_info.id
            iter_start = self.start + worker_id * per_worker
            iter_end = min(iter_start + per_worker, self.end)


        print(f"Worker {worker_info.id if worker_info else 'main'} processing partitions {iter_start} to {iter_end}")
        
        for partition_idx in range(iter_start, iter_end):
            print(self.s3_path)
            partition_path = f"{self.s3_path}/cluster_id={partition_idx}/"

            print(f"Worker {worker_info.id if worker_info else 'main'} processing partition {partition_idx} (path: {partition_path})")
            print(partition_path, "start")

            df_pd = pd.read_parquet(f"{self.s3_path}/cluster_id={partition_idx}/")
            print(partition_path, "ended")
            individual_user_data = df_pd.groupby("idd")

            batch_features = []
            batch_labels = []

            for ii, group in individual_user_data:
                group = group.sort_values(by='date_trans')
                features = group.drop(["idd","id","date_trans", "userid","next_date_trans","cqr"], axis=1).values
                features = torch.tensor(features).float()
                labels = group['cqr'].values
                labels = torch.tensor(labels).float()
                
                num_sequences = len(group) - self.seq_len
            
                for i in range(0, num_sequences):
                    X_sequence = features[i:i+self.seq_len]
                    y_sequence = labels[i+self.seq_len]

                    yield X_sequence, y_sequence.unsqueeze(0)

class PowConDataModule(pl.LightningDataModule):
    def __init__(self, s3_path, s3_path_valid, start, end, seq_len, batch_size, num_workers=1):
        super().__init__()
        self.s3_path = s3_path
        self.s3_path_valid = s3_path_valid
        self.start = start
        self.end = end
        self.seq_len = seq_len
        self.batch_size = batch_size
        self.num_workers = num_workers

    def setup(self, stage=None):
        pass

    def train_dataloader(self):
        return DataLoader(
            PowConDataSetIterable(self.s3_path, "churn_data_train_v4_normalize",self.seq_len, self.batch_size),
            batch_size=512,  # Batch size is controlled by the IterableDataset
            # batch_size = None,
            num_workers=self.num_workers,
            # drop_last=True
        )
    def val_dataloader(self):
        return DataLoader(
            PowConDataSetIterable(self.s3_path_valid, "churn_data_test_high_supply_v4_normalize",self.seq_len, self.batch_size),
            batch_size=512,  # Batch size is controlled by the IterableDataset
            # batch_size = None,
            num_workers=self.num_workers,
            # drop_last=True
        )

class PowConModel(pl.LightningModule):
    
    def __init__(self, n_features, hidden_size, seq_len,
                 batch_size, num_layers, dropout,
                 learning_rate, criterion):
        super().__init__()
        self.n_features = n_features
        self.hidden_size = hidden_size
        self.seq_len = seq_len
        self.batch_size = batch_size
        self.num_layers = num_layers
        self.dropout = dropout
        self.learning_rate = learning_rate
        self.criterion = criterion
        
        self.lstm = nn.LSTM(input_size=n_features,
                            hidden_size=hidden_size,
                            num_layers=num_layers,
                            dropout=dropout,
                            batch_first=True)
        self.batch_norm = nn.BatchNorm1d(hidden_size)
        self.fc = nn.Linear(hidden_size, 1)
        self.sigmoid = nn.Sigmoid()
        self.auroc = BinaryAUROC()
        self.f1score = BinaryF1Score()
        self.acc = BinaryAccuracy()
        
    def forward(self, x):
        lstm_out, _ = self.lstm(x)  # lstm_out = (batch_size, seq_len, hidden_size)
        lstm_out = lstm_out[:, -1, :]
        lstm_out = self.batch_norm(lstm_out)
        x = self.fc(lstm_out)
        x = self.sigmoid(x)
        return x
    
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate, weight_decay=1e-5)
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, patience=3, min_lr=1e-6, verbose=True)
        return {'optimizer': optimizer, 'lr_scheduler': scheduler, 'monitor': 'val_loss'}
    
    
    def training_step(self, train_batch, batch_idx):
        x, y = train_batch
        y_hat = self.forward(x)
        print(x[0][0][-4:])
        # print(y_hat, y)
        loss = self.criterion(y_hat, y)
        self.log('train_loss', loss, prog_bar=True,sync_dist=True)
        return loss

    def validation_step(self, val_batch, batch_idx):
        x, y = val_batch
        y_hat = self.forward(x)
        loss = self.criterion(y_hat, y)
        acc = self.acc(y_hat, y)
        auroc = self.auroc(y_hat, y)
        f1score = self.f1score(y_hat, y)

        self.log('val_loss', loss, prog_bar=True,sync_dist=True)
        self.log('acc', acc, prog_bar=True,sync_dist=True)
        self.log('auroc', auroc, prog_bar=True,sync_dist=True)
        self.log('f1score', f1score, prog_bar=True,sync_dist=True)
        return loss
    

if __name__ == '__main__':
    # Define your parameters
    p = {
        'seq_len': 14,
        'n_features': 15,
        'batch_size': 512,
        'criterion': nn.BCELoss(),
        'max_epochs': 5,
        'hidden_size': 16,
        'num_layers': 5,
        'dropout': 0.5,
        'learning_rate': 0.0001,
        's3_path': "s3://data-science/abcd",
        's3_path_valid': "s3://data-science/abc",
        'n_workers': 2,
        'start': 1,
        'end': 17
    }
    
    # Set random seed for reproducibility
    pl.seed_everything(42)
    
    # Create your data module and model
    data_module = PowConDataModule(
        p['s3_path'], p['s3_path_valid'], p['start'], p['end'], p['seq_len'], p['batch_size'], p['n_workers']
    )
    model = PowConModel(
        n_features=p['n_features'], hidden_size=p['hidden_size'], seq_len=p['seq_len'],
        batch_size=p['batch_size'], criterion=p['criterion'], num_layers=p['num_layers'],
        dropout=p['dropout'], learning_rate=p['learning_rate']
    )
    logger = TensorBoardLogger(save_dir=f"s3://data/logs/{date_time_str}/")
    # Create your trainer and fit your model
    ckpt_path=f"s3://ckpts/{date_time_str}/"
    # Callback to save a checkpoint after each epoch
    checkpoint_every_epoch = ModelCheckpoint(
        dirpath=ckpt_path,
        filename='{epoch:02d}-{val_loss:.2f}',
        save_top_k=-1,  # Save all checkpoints
        verbose=True
    )

    # Callback to save the best checkpoint based on val_loss
    checkpoint_best = ModelCheckpoint(
        dirpath=ckpt_path,
        filename='best-checkpoint',
        save_top_k=1,
        verbose=True,
        monitor='val_loss',
        mode='min'
    )

    lr_monitor = LearningRateMonitor(logging_interval='epoch')

    trainer = pl.Trainer(
        max_epochs=p['max_epochs'],
        accelerator="gpu",
        logger=logger,
        default_root_dir=ckpt_path,
        callbacks=[checkpoint_every_epoch, checkpoint_best, lr_monitor],
        devices=2, strategy="ddp",gradient_clip_val=0.5
    )
    start_time = time.time()
    trainer.fit(model, data_module)
    print(time.time() - start_time)`

Output for reference

2024_07_22_19_28
Global seed set to 42
/home/ray/anaconda3/lib/python3.10/site-packages/fsspec/registry.py:273: UserWarning: Your installed version of s3fs is very old and known to cause
severe performance issues, see also https://github.com/dask/dask/issues/10276

To fix, you should specify a lower version bound on s3fs, or
update the current installation.

  warnings.warn(s3_msg)
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
[rank: 0] Global seed set to 42
Initializing distributed: GLOBAL_RANK: 0, MEMBER: 1/2
2024_07_22_19_28
[rank: 1] Global seed set to 42
/home/ray/anaconda3/lib/python3.10/site-packages/fsspec/registry.py:273: UserWarning: Your installed version of s3fs is very old and known to cause
severe performance issues, see also https://github.com/dask/dask/issues/10276

To fix, you should specify a lower version bound on s3fs, or
update the current installation.

  warnings.warn(s3_msg)
[rank: 1] Global seed set to 42
Initializing distributed: GLOBAL_RANK: 1, MEMBER: 2/2
----------------------------------------------------------------------------------------------------
distributed_backend=nccl
All distributed processes registered. Starting with 2 processes
----------------------------------------------------------------------------------------------------

LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1]
LOCAL_RANK: 1 - CUDA_VISIBLE_DEVICES: [0,1]

  | Name       | Type           | Params
----------------------------------------------
0 | criterion  | BCELoss        | 0     
1 | lstm       | LSTM           | 10.8 K
2 | batch_norm | BatchNorm1d    | 32    
3 | fc         | Linear         | 17    
4 | sigmoid    | Sigmoid        | 0     
5 | auroc      | BinaryAUROC    | 0     
6 | f1score    | BinaryF1Score  | 0     
7 | acc        | BinaryAccuracy | 0     
----------------------------------------------
10.9 K    Trainable params
0         Non-trainable params
10.9 K    Total params
0.043     Total estimated model params size (MB)
Sanity Checking: 0it [00:00, ?it/s]/home/ray/anaconda3/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:224: PossibleUserWarning: The dataloader, val_dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 32 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(
Worker 0 processing partitions 1 to 4
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize
Worker 0 processing partition 1 (path: s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=1/)
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=1/ start
Worker 0 processing partitions 1 to 4
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize
Worker 0 processing partition 1 (path: s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=1/)
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=1/ start
Worker 1 processing partitions 4 to 7
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize
Worker 1 processing partition 4 (path: s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=4/)
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=4/ start
Worker 1 processing partitions 4 to 7
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize
Worker 1 processing partition 4 (path: s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=4/)
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=4/ start
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=4/ ended
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=4/ ended
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=1/ ended
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_test_high_supply_v4_normalize/cluster_id=1/ ended
/home/ray/anaconda3/lib/python3.10/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:224: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 32 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(
Worker 0 processing partitions 1 to 4
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize
Worker 0 processing partition 1 (path: s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=1/)
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=1/ start
Worker 1 processing partitions 4 to 7
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize
Worker 1 processing partition 4 (path: s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=4/)
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=4/ start
Epoch 0: : 0it [00:00, ?it/s]Worker 0 processing partitions 1 to 4
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize
Worker 0 processing partition 1 (path: s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=1/)
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=1/ start
Worker 1 processing partitions 4 to 7
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize
Worker 1 processing partition 4 (path: s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=4/)
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=4/ start
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=1/ ended
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=1/ ended
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=4/ ended
s3://abcd/data-science/akshat-suwalka/xyz_data/xyz_data_train_v4_normalize/cluster_id=4/ ended
tensor([-0.0710,  0.7597, -0.7067,  2.0517], device='cuda:0')
tensor([-0.0710,  0.7597, -0.7067,  2.0517], device='cuda:1')
[rank1]:[W reducer.cpp:1360] Warning: find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters in the forward pass. This flag results in an extra traversal of the autograd graph every iteration,  which can adversely affect performance. If your model indeed never has any unused parameters in the forward pass, consider turning this flag off. Note that this warning may be a false positive if your model has flow control causing later iterations to have unused parameters. (function operator())
[rank0]:[W reducer.cpp:1360] Warning: find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters in the forward pass. This flag results in an extra traversal of the autograd graph every iteration,  which can adversely affect performance. If your model indeed never has any unused parameters in the forward pass, consider turning this flag off. Note that this warning may be a false positive if your model has flow control causing later iterations to have unused parameters. (function operator())
Epoch 0: : 1it [00:09,  9.20s/it, loss=0.746, v_num=0, train_loss=0.746]tensor([-0.3938, -0.1828, -0.5656,  2.3441], device='cuda:1')
tensor([-0.3938, -0.1828, -0.5656,  2.3441], device='cuda:0')
Epoch 0: : 2it [00:09,  4.61s/it, loss=0.747, v_num=0, train_loss=0.748]tensor([-0.3938, -0.1828, -0.2835, -0.6095], device='cuda:0')
tensor([-0.3938, -0.1828, -0.2835, -0.6095], device='cuda:1')
Epoch 0: : 3it [00:09,  3.12s/it, loss=0.742, v_num=0, train_loss=0.732]tensor([-0.3938, -0.1828,  1.4090,  0.4444], device='cuda:1')
tensor([-0.3938, -0.1828,  1.4090,  0.4444], device='cuda:0')
Epoch 0: : 4it [00:09,  2.38s/it, loss=0.738, v_num=0, train_loss=0.727]tensor([-0.3938, -0.1828, -0.7067,  0.1604], device='cuda:0')
tensor([-0.3938, -0.1828, -0.7067,  0.1604], device='cuda:1')
Epoch 0: : 5it [00:09,  1.94s/it, loss=0.737, v_num=0, train_loss=0.733]tensor([-0.3938, -0.1828,  0.8448,  0.6567], device='cuda:1')
tensor([-0.3938, -0.1828,  0.8448,  0.6567], device='cuda:0')
Epoch 0: : 6it [00:09,  1.64s/it, loss=0.738, v_num=0, train_loss=0.739]tensor([ 0.7359, -0.1828, -0.7067,  1.1862], device='cuda:0')
tensor([ 0.7359, -0.1828, -0.7067,  1.1862], device='cuda:1')
Epoch 0: : 7it [00:10,  1.44s/it, loss=0.74, v_num=0, train_loss=0.757] tensor([ 1.8656,  6.4148, -0.7067,  0.8927], device='cuda:1')
tensor([ 1.8656,  6.4148, -0.7067,  0.8927], device='cuda:0')
Epoch 0: : 8it [00:10,  1.27s/it, loss=0.74, v_num=0, train_loss=0.736]tensor([-0.3938, -0.1828, -0.7067,  0.5730], device='cuda:0')
tensor([-0.3938, -0.1828, -0.7067,  0.5730], device='cuda:1')
Epoch 0: : 9it [00:10,  1.16s/it, loss=0.738, v_num=0, train_loss=0.721]tensor([-0.3938, -0.1828, -0.0014,  0.6580], device='cuda:1')
tensor([-0.3938, -0.1828, -0.0014,  0.6580], device='cuda:0')
Epoch 0: : 10it [00:10,  1.05s/it, loss=0.737, v_num=0, train_loss=0.731]tensor([-0.3938, -0.1828,  2.8194, -0.3166], device='cuda:0')
tensor([-0.3938, -0.1828,  2.8194, -0.3166], device='cuda:1')
Epoch 0: : 11it [00:10,  1.02it/s, loss=0.738, v_num=0, train_loss=0.744]tensor([-0.3938, -0.1828,  0.9859, -1.0117], device='cuda:1')
tensor([-0.3938, -0.1828,  0.9859, -1.0117], device='cuda:0')
Epoch 0: : 12it [00:10,  1.11it/s, loss=0.736, v_num=0, train_loss=0.723]tensor([-0.3938, -0.1828, -0.0014, -0.8525], device='cuda:0')
tensor([-0.3938, -0.1828, -0.0014, -0.8525], device='cuda:1')
Epoch 0: : 13it [00:11,  1.17it/s, loss=0.736, v_num=0, train_loss=0.732]tensor([-0.3938, -0.1828, -0.7067, -0.0466], device='cuda:1')
tensor([-0.3938, -0.1828, -0.7067, -0.0466], device='cuda:0')
Epoch 0: : 14it [00:11,  1.25it/s, loss=0.735, v_num=0, train_loss=0.727]tensor([-0.3938, -0.1828,  1.6911, -0.7032], device='cuda:0')
tensor([-0.3938, -0.1828,  1.6911, -0.7032], device='cuda:1')
Epoch 0: : 15it [00:11,  1.31it/s, loss=0.735, v_num=0, train_loss=0.736]tensor([ 1.1125,  2.0164, -0.7067, -0.9280], device='cuda:1')
tensor([ 1.1125,  2.0164, -0.7067, -0.9280], device='cuda:0')
Epoch 0: : 16it [00:11,  1.39it/s, loss=0.736, v_num=0, train_loss=0.741]tensor([-0.3938, -0.1828, -0.2835, -0.7570], device='cuda:0')
tensor([-0.3938, -0.1828, -0.2835, -0.7570], device='cuda:1')
Epoch 0: : 17it [00:11,  1.45it/s, loss=0.734, v_num=0, train_loss=0.706]tensor([-0.3938, -0.1828, -0.8477, -0.7682], device='cuda:1')
tensor([-0.3938, -0.1828, -0.8477, -0.7682], device='cuda:0')
Epoch 0: : 18it [00:11,  1.52it/s, loss=0.733, v_num=0, train_loss=0.716]tensor([-0.3938, -0.1828,  1.2679, -0.9418], device='cuda:0')
tensor([-0.3938, -0.1828,  1.2679, -0.9418], device='cuda:1')
Epoch 0: : 19it [00:12,  1.57it/s, loss=0.732, v_num=0, train_loss=0.710]tensor([-0.3938, -0.1828, -0.5656, -0.7337], device='cuda:1')
tensor([-0.3938, -0.1828, -0.5656, -0.7337], device='cuda:0')
Epoch 0: : 20it [00:12,  1.64it/s, loss=0.732, v_num=0, train_loss=0.740]tensor([-0.3938, -0.1828, -0.7067, -0.5402], device='cuda:0')
tensor([-0.3938, -0.1828, -0.7067, -0.5402], device='cuda:1')
Epoch 0: : 21it [00:12,  1.70it/s, loss=0.731, v_num=0, train_loss=0.722]tensor([-0.3938, -0.1828, -0.5656, -1.4066], device='cuda:1')
tensor([-0.3938, -0.1828, -0.5656, -1.4066], device='cuda:0')
Epoch 0: : 22it [00:12,  1.76it/s, loss=0.729, v_num=0, train_loss=0.710]tensor([-0.3938, -0.1828,  0.8448,  0.8225], device='cuda:0')
tensor([-0.3938, -0.1828,  0.8448,  0.8225], device='cuda:1')
Epoch 0: : 23it [00:12,  1.81it/s, loss=0.729, v_num=0, train_loss=0.730]tensor([ 1.8656, -0.1828, -0.7067, -0.0637], device='cuda:1')
tensor([ 1.8656, -0.1828, -0.7067, -0.0637], device='cuda:0')
Epoch 0: : 24it [00:12,  1.87it/s, loss=0.729, v_num=0, train_loss=0.725]tensor([ 1.8656,  3.1160, -0.7067,  2.5729], device='cuda:0')
tensor([ 1.8656,  3.1160, -0.7067,  2.5729], device='cuda:1')
Epoch 0: : 25it [00:13,  1.92it/s, loss=0.728, v_num=0, train_loss=0.708]tensor([ 1.8656, -0.1828, -0.7067,  2.1329], device='cuda:0')
tensor([ 1.8656, -0.1828, -0.7067,  2.1329], device='cuda:1')
Epoch 0: : 26it [00:13,  1.97it/s, loss=0.727, v_num=0, train_loss=0.732]tensor([ 0.7359,  0.9168, -0.7067, -0.7195], device='cuda:0')
tensor([ 0.7359,  0.9168, -0.7067, -0.7195], device='cuda:1')
Epoch 0: : 27it [00:13,  2.02it/s, loss=0.726, v_num=0, train_loss=0.722]tensor([-0.3938, -0.1828, -0.0014, -1.1270], device='cuda:0')
tensor([-0.3938, -0.1828, -0.0014, -1.1270], device='cuda:1')
Epoch 0: : 28it [00:13,  2.07it/s, loss=0.725, v_num=0, train_loss=0.721]tensor([ 0.5621,  1.0859, -0.7067,  0.5953], device='cuda:0')
tensor([ 0.5621,  1.0859, -0.7067,  0.5953], device='cuda:1')
Epoch 0: : 29it [00:13,  2.12it/s, loss=0.725, v_num=0, train_loss=0.728]tensor([-0.3938, -0.1828,  0.5627,  0.5543], device='cuda:1')
tensor([-0.3938, -0.1828,  0.5627,  0.5543], device='cuda:0')
Epoch 0: : 30it [00:13,  2.17it/s, loss=0.724, v_num=0, train_loss=0.700]tensor([-0.3938, -0.1828, -0.7067,  0.5953], device='cuda:0')
tensor([-0.3938, -0.1828, -0.7067,  0.5953], device='cuda:1')
Epoch 0: : 31it [00:14,  2.21it/s, loss=0.722, v_num=0, train_loss=0.721]tensor([-0.3938, -0.1828, -0.7067, -0.5701], device='cuda:1')
tensor([-0.3938, -0.1828, -0.7067, -0.5701], device='cuda:0')
Epoch 0: : 32it [00:14,  2.25it/s, loss=0.722, v_num=0, train_loss=0.716]tensor([-0.3938, -0.1828,  0.4217, -1.0701], device='cuda:0')