Mihir Patel

Oops meant to open draft PR -- not ready for review yet!

I _think_ this works -- tested on resnet with large batch size, but I'd love to get some help testing this on more workloads. I honestly don't have the bandwidth...

> Is this ready for review? Yes! :)

`I don't understand pyright` same :')

Hm... I'm not getting this on dev... any ideas why?

I'm getting several issues with upgrade... @dblalock are you also seeing these?

@dblalock could you actually just resolve all the lint errors with datasets 2 and then update meta.yaml / setup.py to upgrade to datasets 2?

CC: @moinnadeem do we need any tests for upgrading datasets?

@dblalock I put a bit on your plate, so if you don't have the bandwidth for this, lmk and I'm happy to help out. I'd rather just do the full...

This is going to change following offline discussion