@ptrblck here is my code , i noticed my code is not utilising GPU , though i am assigning device as GPU to model as well as dataset
class SwapDataset:
    def __init__(self, features, noise,batch_size):
        self.features = features
        self.noise = noise
        self.batch_size = batch_size
        
    def __len__(self):
        return (self.features.shape[0])
    
    def __getitem__(self, idx):
        
        sample = self.features[idx, :].copy()
        sample = self.features.copy()
        sample = self.swap_sample(sample)
        
        dct = {
            'x' : torch.tensor(sample[idx, :], dtype=torch.float).to(device),
            'y' : torch.tensor(self.features[idx, :], dtype=torch.float).to(device)
        }
        return dct
    
    def swap_sample(self,sample):
        
        num_samples = self.features.shape[0]
        num_features = self.features.shape[1]

        if len(sample.shape) == 2:
            # batch_size = sample.shape[0]
            random_row = np.random.randint(0, num_samples, size=self.batch_size)

            for i in range(self.batch_size):
                random_col = np.random.rand(num_features) < self.noise
                sample[i,random_col] = self.features[random_row[i], random_col]
        else:
            batch_size = 1
            random_row = np.random.randint(0, num_samples, size=self.batch_size)
            random_col = np.random.rand(num_features) < self.noise
            #sample[:,random_col] = self.features[random_row, random_col]

        return sample


device   = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

train_dataset = SwapDataset(X1,0.02,batch_size=1000) # Swapping data 
data_loader   = torch.utils.data.DataLoader(train_dataset, batch_size=1000, shuffle=False)