Hello, @harsha_g and @SANTOSH_S. I believe the dataset is on CUDA.
Training loop is shown below. Also, args.device=device(type=‘cuda’)
   for _ in train_iterator:
        # train_iterator is basically range(0, total_epochs). However in case we resume training then
        # train_iterator becomes range(completed_epochs, total_epochs)
        # In every epoch we compute the mini-batches again
        train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
        train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)
        epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
         
        for step, batch in enumerate(epoch_iterator):
            # batch is an element in dataloader. i.e. effectively a mini-batch

            # Skip past any already trained steps if resuming training. i.e. number of mini-batches trained already in the epoch
            if steps_trained_in_current_epoch > 0:
                steps_trained_in_current_epoch -= 1
                continue

            model.train()
            batch = tuple(t.to(args.device) for t in batch)

            inputs = {
                "input_ids": batch[0],
                "attention_mask": batch[1],
                "token_type_ids": batch[2],
                "start_positions": batch[3],
                "end_positions": batch[4],
            }

            if args.model_type in ["roberta"]:
                del inputs["token_type_ids"] 
            # roberta doesn't require token_type_ids           

            outputs = model(**inputs)
            # model outputs are always tuple in transformers (see doc)