Skip to content

Instantly share code, notes, and snippets.

@aravindpai

aravindpai/iterator.py

Created January 27, 2020 13:27
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save aravindpai/f8c7d5dfcf50e98dd12c0f6f8aaf4e9b to your computer and use it in GitHub Desktop.
Save aravindpai/f8c7d5dfcf50e98dd12c0f6f8aaf4e9b to your computer and use it in GitHub Desktop.
Download ZIP
Raw
iterator.py
#check whether cuda is available
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#set batch size
BATCH_SIZE = 64
#Load an iterator
train_iterator, valid_iterator = data.BucketIterator.splits(
(train_data, valid_data),
batch_size = BATCH_SIZE,
sort_key = lambda x: len(x.text),
sort_within_batch=True,
device = device)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment