PyTorch to CoreML conversion fails with tools 4.0b1 at BatchNorm2d with track_running_stats=False, ValueError: Input mean is required for op batch_norm.

Title:

PyTorch to CoreML conversion fails with tools 4.0b1 at BatchNorm2d with track_running_stats=False, ValueError: Input mean is required for op batch_norm.

Relevance:

This might seem like an edge case but we have several models where this is relevant. Hence the report.

Documentation at

https://pytorch.org/docs/master/generated/torch.nn.BatchNorm2d.html#batchnorm2d
> If track_running_stats is set to False, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well.

In detail,
affine=True, track_running_stats=True OK
affine=False, track_running_stats=True OK
affine=True, track_running_stats=False FAILS
affine=False, track_running_stats=False FAILS

Reproducible

yes, see testcase attached

[testBN.txt](https://github.com/apple/coremltools/files/4904033/testBN.txt)

Setup

Torch version : 1.5.1
CoreML tools version : 4.0b1
Python 3.7.6
macOS Catalina latest

Log:

Converting Frontend ==> MIL Ops: 91%|██ | 10/11 [00:00<00:00, 4716.94 ops/s]
Traceback (most recent call last):
File "testBN.py", line 29, in
mlmodel = ct.convert( traced_model, inputs= [ ct.TensorType(name="input1", shape=dummy_input.shape) ] )
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/_converters_entry.py", line 299, in convert
**kwargs
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/mil/converter.py", line 120, in _convert
prog = frontend_converter(model, **kwargs)
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/mil/converter.py", line 62, in __call__
return load(*args, **kwargs)
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 86, in load
raise e
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/mil/frontend/torch/load.py", line 76, in load
prog = converter.convert()
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/mil/frontend/torch/converter.py", line 302, in convert
convert_nodes(self.context, self.graph)
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 56, in convert_nodes
_add_op(context, node)
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/mil/frontend/torch/ops.py", line 734, in batch_norm
name=node.name,
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/mil/mil/ops/registry.py", line 62, in add_op
return cls._add_op(op_cls, **kwargs)
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/mil/mil/builder.py", line 186, in _add_op
op_cls.input_spec, kwargs["name"], op_cls, before_op, kwargs
File "~/Library/Python/3.7/lib/python/site-packages/coremltools/converters/mil/mil/builder.py", line 117, in _create_input_vars
"Input {} is required for op {}.".format(in_name, op_cls.__name__)
ValueError: Input mean is required for op batch_norm.