Ensure int64

[knighton]

Description of changes:

tl;dr Fix a crash due to a float creeping into calculations.

shard_shuffle_units = [shard_samples] * (shard_choose // shard_samples)

In the above line of dataset.py, if we are choosing sufficiently few samples out of sufficiently many, the resulting list of shuffle units will be empty (i.e., no shards or partial shards). Numpy concatenate automatically assumes empty lists are of dtype float64. Numpy concatenate also automatically casts lists of both int64 and float64 to float64. This then shows up as the derived quantity num_samples being a float64 in the shuffling algorithms, causing the crash.

Issue #, if available:

Merge Checklist:

Put an x without space in the boxes that apply. If you are unsure about any checklist, please don't hesitate to ask. We are here to help! This is simply a reminder of what we are going to look for before merging your pull request.

General

• I have read the contributor guidelines
• This is a documentation change or typo fix. If so, skip the rest of this checklist.
• I certify that the changes I am introducing will be backward compatible, and I have discussed concerns about this, if any, with the MosaicML team.
• I have updated any necessary documentation, including README and API docs (if appropriate).

Tests

• I ran pre-commit on my change. (check out the pre-commit section of prerequisites)
• I have added tests that prove my fix is effective or that my feature works (if appropriate).
• I ran the tests locally to make sure it pass. (check out testing)
• I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes.