StatefulDataLoader restore error for RNG must be a torch.ByteTensor

Ling_Xue: <p>I save the dataloader as</p>
<pre><code class="lang-auto">from torchdata.stateful_dataloader import StatefulDataLoader

train_dataloader = StatefulDataLoader(
            train_dataset,
            batch_size=args.micro_train_batch_size,
            sampler=None,
            drop_last=True,
            shuffle=True,
            collate_fn=collate_fn,
            pin_memory=False,
        ) 
torch.save({
                    'epoch': epoch, 
                    'model_state_dict': self.markov.state_dict(),
                    'next_sample': global_step,
                    'optimizer_state_dict': self.optimizer_m.state_dict(),  
                    'scheduler_state_dict': self.scheduler_m.state_dict(),  
                    'dataloader_state_dict': self.train_dataloader.state_dict(), 
                    'loss': logs_dict,
                }, save_path)

</code></pre>
<p>I restore the dataloader as follows.</p>
<pre><code class="lang-auto">history = torch.load(save_path, map_location='cuda:0')
dataloader_state = history["dataloader_state_dict"]
</code></pre>
<p>Up till now, there seems to be no problem. But when I use the dataloader to load data, I get the following errors. How can I solve it ?</p>
<p>File “/workspace/generate.py”, line 511, in fit<br>
for inputs, attention_masks, loss_masks in self.train_dataloader:<br>
File “/usr/local/lib/python3.12/dist-packages/torchdata/stateful_dataloader/stateful_dataloader.py”, line 406, in <strong>iter</strong><br>
self._iterator = self._get_iterator()<br>
^^^^^^^^^^^^^^^^^^^^<br>
File “/usr/local/lib/python3.12/dist-packages/torchdata/stateful_dataloader/stateful_dataloader.py”, line 384, in _get_iterator<br>
it = _StatefulSingleProcessDataLoaderIter(self, self.next_iter_state)<br>
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
File “/usr/local/lib/python3.12/dist-packages/torchdata/stateful_dataloader/stateful_dataloader.py”, line 479, in <strong>init</strong><br>
self.load_state_dict(next_iter_state)<br>
File “/usr/local/lib/python3.12/dist-packages/torchdata/stateful_dataloader/stateful_dataloader.py”, line 533, in load_state_dict<br>
self._sampler_iter = try_to_deserialize(self._sampler_iter, state_dict[_SAMPLER_ITER_STATE])<br>
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^<br>
File “/usr/local/lib/python3.12/dist-packages/torchdata/stateful_dataloader/worker.py”, line 57, in try_to_deserialize<br>
obj.load_state_dict(state_dict)<br>
File “/usr/local/lib/python3.12/dist-packages/torchdata/stateful_dataloader/sampler.py”, line 156, in load_state_dict<br>
self.sampler_iter.load_state_dict(state_dict[self._SAMPLER_ITER_STATE])<br>
File “/usr/local/lib/python3.12/dist-packages/torchdata/stateful_dataloader/sampler.py”, line 69, in load_state_dict<br>
self.sampler.generator.set_state(self.generator_state)<br>
TypeError: RNG state must be a torch.ByteTensor</p>

serik-follauter: <p>Hi! So I ran into the very same problem and figured out it was because of setting map_location in torch.load. The RNG state which is supposed to be on the cpu gets sent to cuda and it breaks. In my application I just removed map_location argument, since I run it on the same hardware setup and don’t really need to change devices between saving and loading states. If this is not true for your case you may need to specify map_location more carefully, so that dataloader state gets mapped to cpu and not cuda.</p>