AlexanderLutsenko/nobuco

Add support for multidimensional aggregating function

alievrusik opened this issue · 1 comments

First of all thank you for your amazing job on this project, it comes so handy for my projects.

Secondly, I found out that current default math conversion operations like sum, mean, etc. does not support multidimensional aggregation. For example, if pytorch code contains something like my_tensor.mean((2,3)), nobuco throws error during conversion.

TypeError: '<' not supported between instances of 'tuple' and 'int'

Code to reproduce:

class DummyModel(nn.Module):
    
    def __init__(self):
        super().__init__()
        
    def forward(self, x):
        return x.mean((2,3), keepdim=True)

model = DummyModel()

dummy_image = torch.randn(1, 3, 100, 100)

keras_model = nobuco.pytorch_to_keras(
    model,
    args=[dummy_image]
)