How to work with two input or more about flops-counter.pytorch HOT 6 CLOSED

import torch
import torch.nn as nn
from ptflops import get_model_complexity_info


class Siamese(nn.Module):
    def __init__(self):
        super(Siamese, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, 3, 1)
        self.conv2 = nn.Conv2d(1, 10, 3, 1)

    def forward(self, x1, x2):
        return self.conv1(x1) + self.conv2(x2)

def prepare_input(resolution):
    x1 = torch.FloatTensor(1, *resolution)
    x2 = torch.FloatTensor(1, *resolution)
    return dict(x1 = x1), dict(x2 = x2)

if __name__ == '__main__':
    model = Siamese()
    flops, params = get_model_complexity_info(model, input_res=(1, 224, 224),
                                              input_constructor=prepare_input,
                                              as_strings=True, print_per_layer_stat=False)
    print('      - Flops:  ' + flops)
    print('      - Params: ' + params)