v0.0.4
fn (%data: Tensor[(5, 512, 14, 14), int32],
    %weight: Tensor[(512, 512, 3, 3), int32]
    ) -> Tensor[(5, 512, 14, 14), int32] {
  let %x: Tensor[(5, 512, 14, 14), int32] =
  nn.conv2d(
    %data,
    %weight,
    padding=[1, 1, 1, 1],
    channels=512,
    kernel_size=[3, 3]
  );
  %x
}