Commit 7388f508 authored by Tal Sh's avatar Tal Sh Committed by Yuxin Wu

Fix RepeatedData + PrefetchDataZMQ (#747)

Doing:

```
ds = RepeatedData(ds, -1)
ds = PrefetchDataZMQ(ds, num_workers)
ds.reset_state()
```

Would blow up, since `RepeatedData` raises a ValueError, and `PrefetchDataZMQ` only checks for `NotImplementedError`.

Not the cleanest fix but 🤷‍♂️
parent e6bd7188
...@@ -330,7 +330,7 @@ class RepeatedData(ProxyDataFlow): ...@@ -330,7 +330,7 @@ class RepeatedData(ProxyDataFlow):
:class:`ValueError` when nr == -1. :class:`ValueError` when nr == -1.
""" """
if self.nr == -1: if self.nr == -1:
raise ValueError("size() is unavailable for infinite dataflow") raise NotImplementedError("size() is unavailable for infinite dataflow")
return self.ds.size() * self.nr return self.ds.size() * self.nr
def get_data(self): def get_data(self):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment