tf_efficientnet_lite ONNX export

Hi again,

As per my previous issue, efficientnet_b0 weights exported to ONNX perfectly. But I went to try tf_efficientnet_lite2 today and I have encountered a similar issue to #34 . However the thing which is confusing me is that it works on Colab and not on my local machine with the a virtual environment with the exact same torch 1.5.0+cu101 packaged installed via pip (as opposed to conda for the sake of comparison) and exact same code. Could you please point me in the right direction to fix this problem? Thank you!

Here is the Colab notebook for reference: https://colab.research.google.com/drive/1jG7mSPitb-acA7EgMFaMsG3u57c51Bbe

And this what my local machine is telling me each time I run the exact same code.

(venv) andrewl@albona:/albona/nobackup/andrewl/test$ python3 test.py 
/pytorch/aten/src/ATen/native/BinaryOps.cpp:81: UserWarning: Integer division of tensors using div or / is deprecated, and in a future release div will perform true division as in Python 3. Use true_divide or floor_divide (// in Python) instead.
/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/geffnet/conv2d_layers.py:39: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  return max((math.ceil(i / s) - 1) * s + (k - 1) * d + 1 - i, 0)
/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/geffnet/conv2d_layers.py:39: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  return max((math.ceil(i / s) - 1) * s + (k - 1) * d + 1 - i, 0)
Traceback (most recent call last):
  File "test.py", line 23, in <module>
    output_names=['output'],
  File "/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/torch/onnx/__init__.py", line 168, in export
    custom_opsets, enable_onnx_checker, use_external_data_format)
  File "/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/torch/onnx/utils.py", line 69, in export
    use_external_data_format=use_external_data_format)
  File "/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/torch/onnx/utils.py", line 488, in _export
    fixed_batch_size=fixed_batch_size)
  File "/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/torch/onnx/utils.py", line 351, in _model_to_graph
    fixed_batch_size=fixed_batch_size, params_dict=params_dict)
  File "/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/torch/onnx/utils.py", line 154, in _optimize_graph
    graph = torch._C._jit_pass_onnx(graph, operator_export_type)
  File "/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/torch/onnx/__init__.py", line 199, in _run_symbolic_function
    return utils._run_symbolic_function(*args, **kwargs)
  File "/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/torch/onnx/utils.py", line 740, in _run_symbolic_function
    return op_fn(g, *inputs, **attrs)
  File "/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py", line 128, in wrapper
    args = [_parse_arg(arg, arg_desc) for arg, arg_desc in zip(args, arg_descriptors)]
  File "/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py", line 128, in <listcomp>
    args = [_parse_arg(arg, arg_desc) for arg, arg_desc in zip(args, arg_descriptors)]
  File "/nb/andrewl/anaconda3/envs/venv/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py", line 81, in _parse_arg
    "', since it's not constant, please try to make "
RuntimeError: Failed to export an ONNX attribute 'onnx::Cast', since it's not constant, please try to make things (e.g., kernel size) static if possible

I've also collected my system information for you here:

Collecting environment information...
PyTorch version: 1.5.0+cu101
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Debian GNU/Linux 10 (buster)
GCC version: (Debian 8.3.0-6) 8.3.0
CMake version: Could not collect

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 9.2.148
GPU models and configuration: 
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti

Nvidia driver version: 418.74
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.16.2
[conda] numpy                     1.18.3                   pypi_0    pypi
[conda] torch                     1.5.0+cu101              pypi_0    pypi
[conda] torchvision               0.6.0+cu101              pypi_0    pypi

And the information from the Colab notebook:

Collecting environment information...
PyTorch version: 1.5.0+cu101
Is debug build: No
CUDA used to build PyTorch: 10.1

OS: Ubuntu 18.04.3 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.12.0

Python version: 3.6
Is CUDA available: Yes
CUDA runtime version: 10.1.243
GPU models and configuration: GPU 0: Tesla P100-PCIE-16GB
Nvidia driver version: 418.67
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

Versions of relevant libraries:
[pip3] numpy==1.18.3
[pip3] torch==1.5.0+cu101
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.3.1
[pip3] torchvision==0.6.0+cu101
[conda] Could not collect

rwightman / gen-efficientnet-pytorch

tf_efficientnet_lite ONNX export #39