search

silvandeleemput / memcnn

PyTorch Framework for Developing Memory Efficient Deep Invertible Networks
MIT License
251 stars 26 forks source link

Custom split dimension #49

Closed ClashLuke closed 4 years ago

ClashLuke commented 4 years ago

Description

It'd be great to be able to split at dimensions other than 1, enabling the usage of other layers. Currently MemCNN would split the 1st dimension, which is the batch or sequence dimension in case of RNNs, at all times.

silvandeleemput commented 4 years ago

@ClashLuke this would indeed be a useful feature. It should be relatively easy to add to the AdditiveCoupling and AffineCoupling classes by adding a split_dim attribute. I also only want to add it to default implementations (-1) since the rest are deprecated and will be removed in the future. I'll try to implement this feature soon.

TODO