fix: update TinyStories optimize to use item_loader=TokensLoader()

[bhimrazy]

What does this pr do ?

This PR updates TinyStories optimize to use item_loader=TokensLoader() as per the rule from LitData, i.e.:

outputs = optimize(
    ...,
    # This is important to inform LitData that we are encoding a contiguous 1D array (tokens). 
    # LitData skips storing metadata for each sample, e.g., all the tokens are concatenated to form one large tensor.
    item_loader=TokensLoader(),
)

It was addressed in #2048 by @andyland , but the change was limited to the test and seems to be missed in the source.

Fixes #2144

Additional Info

Screenshots

Before: failing with error

After: