I am training a set of FF networks on tabular data.  The input is a sparse matrix.  Here’s the relevant code:
    #training
    self.data_tr = TensorDataset(
            torch.tensor(train_csc.toarray(), dtype=torch.float32, device=self.device),
            torch.tensor(train_pd['is_case'].values, dtype=torch.float32, device=self.device)   #labels
        )
    
    #validation
    self.data_va = TensorDataset(
            torch.tensor(valid_csc.toarray(), dtype=torch.float32, device=self.device),
            torch.tensor(valid_pd['is_case'].values, dtype=torch.float32, device=self.device)   #labels
        )

and used it in the training
    train_ldr = DataLoader(dataset=self.data_tr, batch_size=param['bs'], shuffle=True)
    for X_mb,y_mb in train_ldr:
        yhat_mb = model(X_mb)
        loss = criterion(yhat_mb[:,0], y_mb)
        ...

The dense array is being stored on the GPU and sliced as required.  This runs very fast.  Unfortunately, a couple of instances are so big that they do not fit on the GPU memory as required in the above approach.  For those instances I have the following:
class SparseDataset(Dataset):

    def __init__(self, mat_csc, label, device='cpu'):
        self.dim = mat_csc.shape
        self.device = torch.device(device)

        csr = mat_csc.tocsr(copy=True)
        self.indptr = torch.tensor(csr.indptr, dtype=torch.int64, device=self.device)
        self.indices = torch.tensor(csr.indices, dtype=torch.int64, device=self.device)
        self.data = torch.tensor(csr.data, dtype=torch.float32, device=self.device)

        self.label = torch.tensor(label, dtype=torch.float32, device=self.device)

    def __len__(self):
        return self.dim[0]

    def __getitem__(self, idx):
        obs = torch.zeros((self.dim[1],), dtype=torch.float32, device=self.device)
        ind1,ind2 = self.indptr[idx],self.indptr[idx+1]
        obs[self.indices[ind1:ind2]] = self.data[ind1:ind2]

        return obs,self.label[idx]

instantiated as
    self.data_tr = SparseDataset(train_csc, train_pd['is_case'].values, device)
    self.data_va = SparseDataset(valid_csc, valid_pd['is_case'].values, device)

and used as
    train_ldr = DataLoader(dataset=self.data_tr, batch_size=param['bs'], shuffle=True, collate_fn=my_collate)
    for X_mb,y_mb in train_ldr:
        yhat_mb = model(X_mb)
        loss = criterion(yhat_mb[:,0], y_mb)
        ....

While this is VERY memory efficient, even on my smallest instance (which fits in the memory), it is 20 times slower than the first approach.  I am looking for ideas to make this faster.  Thx.