Doesn't built on Python 3.9 / PyTorch master due to wrong type annotations

[t-vi]

🐛 Bug

__iter__ in pl_bolts/datamodules/experience_source.py is annotated as returning an Iterable when it should return an Iterator (__iter__ is implemented on Iterables to give Iterators).

To Reproduce

With lightning bolts master and Python 3.9 and PyTorch master:

python3 setup.py  bdist_wheel
Traceback (most recent call last):
  File "/home/tv/pytorch/lightning/lightning-bolts/setup.py", line 13, in <module>
    from pl_bolts import __about__ as about
  File "/home/tv/pytorch/lightning/lightning-bolts/pl_bolts/__init__.py", line 11, in <module>
    from pl_bolts import (  # noqa: E402
  File "/home/tv/pytorch/lightning/lightning-bolts/pl_bolts/datamodules/__init__.py", line 5, in <module>
    from pl_bolts.datamodules.experience_source import DiscountedExperienceSource, ExperienceSource, ExperienceSourceDataset
  File "/home/tv/pytorch/lightning/lightning-bolts/pl_bolts/datamodules/experience_source.py", line 24, in <module>
    class ExperienceSourceDataset(IterableDataset):
  File "/usr/local/lib/python3.9/dist-packages/torch/utils/data/_typing.py", line 256, in __new__
    return super().__new__(cls, name, bases, namespace)
  File "/usr/lib/python3.9/abc.py", line 85, in __new__
    cls = super().__new__(mcls, name, bases, namespace, **kwargs)
  File "/usr/local/lib/python3.9/dist-packages/torch/utils/data/_typing.py", line 330, in _dp_init_subclass
    raise TypeError("Expected 'Iterator' as the return annotation for `__iter__` of {}"

Expected behavior

Working setup.py

Environment

• PyTorch 1.9 (master)
• OS (e.g., Linux): Linux
• How you installed PyTorch (conda, pip, source): self-built + pip

Additional context

I didn't check more type annotations.

[github-actions]

Hi! thanks for your contribution!, great first issue!

[amogkam]

Can a new release be made with this fix? Otherwise lightning-bolts is unusable with torch 1.9 which got released yesterday.