DataLoader Not Iterating Correctly

slyle: <p>I have a pandas dataframe containing tokenized input sequences and labels in list format.</p>
<p>However, I’m having a hard time creating a CustomDataset and DataLoader to convert these lists to torch tensors and batch them together. I’ve created a dummy pandas dataframe for demonstration.  Here is my code:</p>
<pre><code class="lang-auto">#dummy dataframe
import random
data = {
    "text_input_ids": [[random.randint(1, 100) for _ in range(10)] for _ in range(100)],  # Lists of length 10
    "d_icd_desc_code_label": [[random.randint(0, 1) for _ in range(15)] for _ in range(100)],  # Lists of length 15
}

# Create the DataFrame
test_df = pd.DataFrame(data)

class CustomDataset(Dataset):
    def __init__(self, dataframe):
        assert isinstance(dataframe, pd.DataFrame), 'dataframe must be of type pd.DataFrame'
        self.dataframe = dataframe

    def __len__(self):
        return len(self.dataframe)

    def __getitem__(self, idx):
        input_tokens = self.dataframe.iloc[idx]['text_input_ids'] 
        label_tokens = self.dataframe.iloc[idx]['d_icd_desc_code_label']
        return torch.tensor(input_tokens), torch.tensor(label_tokens)

train_ds = CustomDataset(dataframe=test_df)
</code></pre>
<p>Checking one element in the Dataset object:</p>
<pre><code class="lang-auto">train_ds[0]
</code></pre>
<p>returns:</p>
<pre><code class="lang-auto">(tensor([15, 56, 23,  8, 45, 23, 36, 20, 86, 21]),
 tensor([0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 1]))
</code></pre>
<p>However, attempting:</p>
<pre><code class="lang-auto">train_loader = torch.utils.data.DataLoader(train_ds,batch_size=2,shuffle=True)
next(iter(train_loader))
</code></pre>
<p>returns</p>
<pre><code class="lang-auto">Cell In[803], line 12, in CustomDataset.__getitem__(self, idx)
     10 input_tokens = self.dataframe.iloc[idx]['text_input_ids'] 
     11 label_tokens = self.dataframe.iloc[idx]['d_icd_desc_code_label']
---&gt; 12 return torch.tensor(input_tokens), torch.tensor(label_tokens)

ValueError: could not determine the shape of object type 'Series'
</code></pre>
<p>and attempting</p>
<pre><code class="lang-auto">train_loader = torch.utils.data.DataLoader(train_ds,batch_size=2,shuffle=False)
next(iter(train_loader))
</code></pre>
<p>returns</p>
<pre><code class="lang-auto">File c:\Users\SPADMINKH\Documents\Repo\Auto Coding MIMIC Prototype\venv-autocoding\Lib\site-packages\datasets\arrow_dataset.py:2808, in Dataset.__getitems__(self, keys)
   2806 """Can be used to get a batch using a list of integers indices."""
   2807 batch = self.__getitem__(keys)
-&gt; 2808 n_examples = len(batch[next(iter(batch))])
   2809 return [{col: array[i] for col, array in batch.items()} for i in range(n_examples)]

TypeError: only integer tensors of a single element can be converted to an index
</code></pre>
<p>I’ve been at this for hours!  Can someone help me identify my error here?</p>

slyle: <p>I am really, and I mean REALLY dense.</p>
<p>There was no error in my code.  I was importing huggingface datasets Dataset instead of importing the torch Dataset class, so it was obviously inheriting the wrong class Dataset when I instantiated my data loader.</p>
<p>Good grief!  Don’t be like me!</p>

ptrblck: <p>That’s interesting to hear as I also executed your code locally before reading your update here and it works fine using the <code>Dataset</code> class from PyTorch.<br>
The posted error messages you were seeing are also quite confusing.</p>