I see. Yeah I could look at that but it needs to be re-implemented in pytorch to provide the backwards pass anyway. But the implementation might be a useful resource.