Skip to content

Instantly share code, notes, and snippets.

@mirth

mirth/squeeze_collate.py

Last active August 23, 2020 21:50
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 mirth/835efa5c50858874f94db07fcef3ee29 to your computer and use it in GitHub Desktop.
Save mirth/835efa5c50858874f94db07fcef3ee29 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
squeeze_collate.py
def squeeze_collate(batch):
batch_x, batch_y = list(zip(*batch))
batch_y = np.array(batch_y, dtype=np.float32)
batch_y = torch.from_numpy(batch_y)
batch_x = torch.cat(batch_x, dim=0)
return batch_x, batch_y
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment