Improve more on new dataset API

[zasdfgbnm]

The new suggested way to use dataset is:

training, validation = torchani.data.load(dspath).subtract_self_energies(energy_shifter).species_to_indices().shuffle().split(0.8, None)
training = training.collate(batch_size).cache()
validation = validation.collate(batch_size).cache()

which is cleaner than the previous

dataset = torchani.data.load(dspath).subtract_self_energies(energy_shifter).species_to_indices().shuffle()
size = len(dataset)
training, validation = dataset.split(int(0.8 * size), None)
training = training.collate(batch_size).cache()
validation = validation.collate(batch_size).cache()

The memory usage has reduced from the previous 18GB to 16GB, this is a small improvement, but it is hard to make it as good as the previously API by @yueyericardo because we are now not chunking, so we have to spend lots of memories to store padding.

If we don't want to use that much memory, we can remove the .cache() from the last two lines, but we now need a data loader to achieve comparable performance:

training, validation = torchani.data.load(dspath).subtract_self_energies(energy_shifter).species_to_indices().shuffle().split(0.8, None)
training = torch.utils.data.DataLoader(list(training), batch_size=batch_size, collate_fn=torchani.data.collate_fn, num_workers=64)
validation = torch.utils.data.DataLoader(list(validation), batch_size=batch_size, collate_fn=torchani.data.collate_fn, num_workers=64)

The above code uses ~9GB RAM

[farhadrgh]

LGTM