Iterable Datasets work with TPUs #17321
Labels
accelerator: tpu
Tensor Processing Unit
feature
Is an improvement or enhancement
help wanted
Open to be worked on
Milestone
Bug description
The error
TPUs do not currently support IterableDataset objects, the dataset must implement
len. HINT: You can mock the length on your dataset to bypass this error.
is raised when using IterableDatasets on TPUs (code location for fabric and pytorch). At the time this error was introduced (#6875), there was a line in the xla code that would call forlen()
. However, this issue from the xla side was fixed shortly after in pytorch/xla#2866. #6955 was then submitted to request a fix for the error check, but was never resolved. Since #6955 is from two years ago and I am unable to open the issue, I wanted to revive this topic in a new issue.For reference, I ran into this warning when running the dqn lightning tutorial on TPU(s).
Unless there is a good reason to keep this error check, can we either remove it entirely or adjust it to a warning saying something along the lines of "If you want to use
len()
on your IterableDataset, you will need to manually implementlen()
?What version are you seeing the problem on?
1.6.x, 1.7.x, 1.8.x, 1.9.x, 2.0+
cc @Borda @JackCaoG @steventk-g @Liyang90 @carmocca
The text was updated successfully, but these errors were encountered: