optimize_for_mobile() assert fail: missing op for prepacked::conv2d_clamp_prepack

[BryceEakin]

🐛 Bug

When calling torch.utils.mobile_optimizer.optimize_for_mobile, an internal assert fails:

0INTERNAL ASSERT FAILED at "../torch/csrc/jit/ir/alias_analysis.cpp":532, please report a bug to PyTorch. We don't have an op for prepacked::conv2d_clamp_prepack but it isn't a special case.  Argument types: Tensor, Tensor, int[], str, int[], int, NoneType, NoneType,

To Reproduce

Steps to reproduce the behavior:

Created convolutional model (gapnet plus linear -> sigmoid head). Resulting structure as output by pytorch:

FullGuidanceModel(
  (plan_model): GuidancePlanModel(
    (backbone): GapSmall(
      (model): Sequential(
        (0): TensorflowSamePadForConv2D()
        (1): Conv2d(1, 16, kernel_size=(3, 3), stride=(2, 2), padding=valid)
        (2): ELU(alpha=1.0)
        (3): TensorflowSamePadForConv2D()
        (4): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2), padding=valid)
        (5): ELU(alpha=1.0)
        (6): TensorflowSamePadForConv2D()
        (7): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2), padding=valid)
        (8): ELU(alpha=1.0)
        (9): TensorflowSamePadForConv2D()
        (10): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=valid)
        (11): ELU(alpha=1.0)
        (12): TensorflowSamePadForConv2D()
        (13): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=valid)
        (14): ELU(alpha=1.0)
      )
    )
    (logits_dense): Linear(in_features=128, out_features=1, bias=True)
  )
  (moving_avg_layer): MovingAverageAcrossBatchWithMemory()
)

Note: TensorflowSamePadForConv2D modules just apply a pre-calculated F.pad() operation to their input and return. GapSmall's output gets averaged over spatial dimensions.

Expected behavior

scripted module compiles normally.

Environment

• PyTorch Version (e.g., 1.0): 1.9.1+cu102
• OS (e.g., Linux): Ubuntu 20.04 (WSL)
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source):
• Python version: 3.8.7
• CUDA/cuDNN version: n/a
• GPU models and configuration: n/a
• Any other relevant information:

Additional context

[linbinyu]

@kimishpatel any idea?

[kimishpatel]

@BryceEakin did you build from source and if so did you set USE_XNNPACK=0?

[zrfisher]

I'm also receiving this error on MacOS 12.1 (Intel Core i7) and when running on Google Colab