Torch script export not working with hugging face model

[HassanFahmy]

When trying to export a model that was built using the hugging face class and trained with composer to torchscript a line of code in forward function in the hf_from_composer_checkpoint function causes the following error.

https://github.com/mosaicml/composer/blob/6f18ff8641bd87e13f0d51b3b152e4335e6a4b28/composer/models/huggingface.py#L488C1-L489C1

W0827 18:15:51.775183 140737350283712 inference.py:225] Scripting with torch.jit.script failed and sample inputs are not provided for tracing with torch.jit.trace
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/composer/utils/inference.py", line 219, in export_for_inference
    export_model = torch.jit.script(model)
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/_script.py", line 1338, in script
    return torch.jit._recursive.create_script_module(
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/_recursive.py", line 558, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/_recursive.py", line 571, in create_script_module_impl
    method_stubs = stubs_fn(nn_module)
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/_recursive.py", line 898, in infer_methods_to_compile
    stubs.append(make_stub_from_method(nn_module, method))
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/_recursive.py", line 87, in make_stub_from_method
    return make_stub(func, method_name)
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/_recursive.py", line 71, in make_stub
    ast = get_jit_def(func, name, self_name="RecursiveScriptModule")
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/frontend.py", line 372, in get_jit_def
    return build_def(
  File "/usr/local/lib/python3.10/dist-packages/torch/jit/frontend.py", line 430, in build_def
    type_comment_decl = torch._C.parse_type_comment(type_line)
RuntimeError: expected type comment but found 'ident' here:
output = self.model(**batch)  # type: ignore (thirdparty)
~~~~~~ <--- HERE

code to repro:

#model = transformers.AutoModelForImageClassification.from_pretrained("google/efficientnet-b4")
 model = composer.models.HuggingFaceModel(model)
 model.load_state_dict(
        torch.load(checkpoint_path_abs, map_location='cuda')["state"]["model"]
    )
 export_for_inference(model=model,
                     save_format='torchscript',
                     save_path=os.path.join(args.export_dir_path, 'model.pt'))

[cli99]

Removing the comment at output = self.model(**batch)gives

    return method(ctx, node)
           ^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3/dist-packages/torch/jit/frontend.py", line 1245, in build_DictComp
    raise NotSupportedError(r, "Comprehension ifs are not supported yet")
torch.jit.frontend.NotSupportedError: Comprehension ifs are not supported yet:
  File "/usr/lib/python3/dist-packages/composer/models/huggingface.py", line 487
        if isinstance(batch, Mapping):
            # Further input validation is left to the huggingface forward call
            batch = {k: v for k, v in batch.items() if k in self.model_forward_args}
            output = self.model(**batch)
        else:

save_format='torchscript' does not seem to be supported by torch/HF yet for the model.