code: def test_flop_count_conv(
      input_shape,
      layer_channels,
      layer_strides,
      layer_kernels,
      layer_padding,
  ):
      # Testing that the flop counter counts a convolutional model as expected
      # (i.e. when compared to a reference value).
  
      class ConvModel(torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.layers = torch.nn.ModuleList()
              cin = input_shape[1]
  
              for cout, stride, kernel, padding in zip(
                  layer_channels,
                  layer_strides,
                  layer_kernels,
                  layer_padding,
              ):
                  self.layers.append(
                      torch.nn.Conv2d(
                          cin,
                          cout,
                          kernel_size=kernel,
                          stride=stride,
                          padding=padding,
                      )
                  )
                  cin = cout
  
          def forward(self, inp):
              flops = 0
              out = inp
              for layer, k_size in zip(self.layers, layer_kernels):
                  cin = out.shape[1]
                  out = layer(out)
                  # +3 for counting the additions to the flops variable itself
                  flops += 3 + out.numel() * k_size * k_size * cin + 3
  
              return out, flops
  
      model = ConvModel()
      inp = torch.randn(input_shape)
      _, correct_flops = model(inp)
  
      counted_flops, _, unsupported_ops = count_flops(model, (inp,))
      assert counted_flops == correct_flops
      assert len(unsupported_ops) == 0
  
imports: import pytest
  import torch
  import torch.nn
  from neuralcompression.functional import count_flops
label_desc: Use image node_id 2 for calling a global function with example usage: test_flop_count_conv(input_shape, layer_channels, layer_strides, layer_kernels, layer_padding) and returns: out,  flops