Getting KeyError while trying to iterate through a dataset with Dataloader

ggg: <p>Hi!</p>
<p>I have a dataframe called <strong>df</strong> with the columns called <strong>word</strong> and <strong>dict</strong>.</p>
<p>Firstly, I split this dataframe into train and test parts with help of <strong>train_test_split</strong>.</p>
<pre><code class="lang-auto">train_df, test_df = train_test_split(df, test_size=0.2, shuffle=True)
</code></pre>
<p>Then I created a custom dataset, which I then passed into a Dataloader, so as to iterate through the data in the train function of the machine learning model.</p>
<pre><code class="lang-auto">class MyDataset(Dataset):
    
    def __init__(self, df):       
        self.words = df['word']
        self.dicts = df['dict']
    
    def __len__(self):
        return len(self.words)
    
    def __getitem__(self, index):
        word_ = self.words[index]
        dict_ = self.dicts(index)
        
        return word_, dict_
</code></pre>
<pre><code class="lang-auto">train_dataset = MyDataset(train_df)
test_dataset = MyDataset(test_df)
train_loader = DataLoader(train_dataset, batch_size=10)
test_loader = DataLoader(test_dataset, batch_size=10)
</code></pre>
<p>The problem is that if I try to iterate through MyDataset I keep getting KeyErrors. It seems like the Dataloader tries to iterate through indexes of test_dataset, when it goes through train and vice versa. I tried to look for the solution of this problem, but have not found it anywhere.<br>
Do you have an idea, how I can fix it?<br>
Thank you in advance for your help!</p>
<p>The train_dataloader loop looks as follows (just for testing):</p>
<pre><code class="lang-auto">for x, y in train_loader:
    print(x, y)
</code></pre>
<p>The error I get is this:</p>
<pre><code class="lang-auto">KeyError                                  Traceback (most recent call last)
&lt;ipython-input-93-dcdc8a11e4d8&gt; in &lt;module&gt;
----&gt; 1 for x, y in train_loader:
      2     print(x, y)

C:\Downloads\Anaconda\lib\site-packages\torch\utils\data\dataloader.py in __next__(self)
    433         if self._sampler_iter is None:
    434             self._reset()
--&gt; 435         data = self._next_data()
    436         self._num_yielded += 1
    437         if self._dataset_kind == _DatasetKind.Iterable and \

C:\Downloads\Anaconda\lib\site-packages\torch\utils\data\dataloader.py in _next_data(self)
    473     def _next_data(self):
    474         index = self._next_index()  # may raise StopIteration
--&gt; 475         data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    476         if self._pin_memory:
    477             data = _utils.pin_memory.pin_memory(data)

C:\Downloads\Anaconda\lib\site-packages\torch\utils\data\_utils\fetch.py in fetch(self, possibly_batched_index)
     42     def fetch(self, possibly_batched_index):
     43         if self.auto_collation:
---&gt; 44             data = [self.dataset[idx] for idx in possibly_batched_index]
     45         else:
     46             data = self.dataset[possibly_batched_index]

C:\Downloads\Anaconda\lib\site-packages\torch\utils\data\_utils\fetch.py in &lt;listcomp&gt;(.0)
     42     def fetch(self, possibly_batched_index):
     43         if self.auto_collation:
---&gt; 44             data = [self.dataset[idx] for idx in possibly_batched_index]
     45         else:
     46             data = self.dataset[possibly_batched_index]

C:\Downloads\Anaconda\lib\site-packages\pandas\core\frame.py in __getitem__(self, key)
   2798             if self.columns.nlevels &gt; 1:
   2799                 return self._getitem_multilevel(key)
-&gt; 2800             indexer = self.columns.get_loc(key)
   2801             if is_integer(indexer):
   2802                 indexer = [indexer]

C:\Downloads\Anaconda\lib\site-packages\pandas\core\indexes\base.py in get_loc(self, key, method, tolerance)
   2646                 return self._engine.get_loc(key)
   2647             except KeyError:
-&gt; 2648                 return self._engine.get_loc(self._maybe_cast_indexer(key))
   2649         indexer = self.get_indexer([key], method=method, tolerance=tolerance)
   2650         if indexer.ndim &gt; 1 or indexer.size &gt; 1:

pandas\_libs\index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas\_libs\index.pyx in pandas._libs.index.IndexEngine.get_loc()

pandas\_libs\hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

pandas\_libs\hashtable_class_helper.pxi in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 0
</code></pre>

Dwight_Foster: <p>The problem is that you cannot index a pandas dataframe. You should be able to do this instead</p>
<pre><code class="lang-auto">self.words.iloc[index]
</code></pre>

ggg: <p>Thanks for your answer! It solved the problem.</p>