[Bug]: CS-Flow torch model export fails with wandb version >0.15.9

[ashwinvaidya17]

Describe the bug

Need to revisit the dependency. CS-Flow torch model export fails with wandb version >0.15.9.

Dataset

N/A

Model

CSFlow

Steps to reproduce the behavior

1. Install wandb>=0.5.10 on your environment.
2. Call torch export on a trained CS-Flow model

anomalib export --model Csflow --export_type TORCH --ckpt_path tmp/Csflow/MVTec/dummy/latest/weights/lightning/model.ckpt --export_root export/new

OS information

OS information:

• OS: [e.g. Ubuntu 20.04]: Ubuntu 20.04
• Python version: [e.g. 3.10.0]: 3.10.4
• Anomalib version: [e.g. 0.3.6]: main branch
• PyTorch version: [e.g. 1.9.0]: 2.1.2+cu11
• CUDA/cuDNN version: [e.g. 11.1]: 11
• GPU models and configuration: [e.g. 2x GeForce RTX 3090]: 3090
• Any other relevant information: [e.g. I'm using a custom dataset]

Expected behavior

Torch export should work

Screenshots

No response

Pip/GitHub

GitHub

What version/branch did you use?

main

Configuration YAML

NA

Logs

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /home/ashwin/miniconda3/envs/temp/bin/anomalib:8 in <module>                                     │
│                                                                                                  │
│   5 from anomalib.cli.cli import main                                                            │
│   6 if __name__ == '__main__':                                                                   │
│   7 │   sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])                         │
│ ❱ 8 │   sys.exit(main())                                                                         │
│   9                                                                                              │
│                                                                                                  │
│ /home/ashwin/projects/anomalib_copy/src/anomalib/cli/cli.py:470 in main                          │
│                                                                                                  │
│   467 def main() -> None:                                                                        │
│   468 │   """Trainer via Anomalib CLI."""                                                        │
│   469 │   configure_logger()                                                                     │
│ ❱ 470 │   AnomalibCLI()                                                                          │
│   471                                                                                            │
│   472                                                                                            │
│   473 if __name__ == "__main__":                                                                 │
│                                                                                                  │
│ /home/ashwin/projects/anomalib_copy/src/anomalib/cli/cli.py:63 in __init__                       │
│                                                                                                  │
│    60 │   │   if _LIGHTNING_AVAILABLE:                                                           │
│    61 │   │   │   self.before_instantiate_classes()                                              │
│    62 │   │   │   self.instantiate_classes()                                                     │
│ ❱  63 │   │   self._run_subcommand()                                                             │
│    64 │                                                                                          │
│    65 │   def init_parser(self, **kwargs) -> ArgumentParser:                                     │
│    66 │   │   """Method that instantiates the argument parser."""                                │
│                                                                                                  │
│ /home/ashwin/projects/anomalib_copy/src/anomalib/cli/cli.py:351 in _run_subcommand               │
│                                                                                                  │
│   348 │   │   elif self.config["subcommand"] in (*self.subcommands(), "train", "export", "pred   │
│   349 │   │   │   fn = getattr(self.engine, self.subcommand)                                     │
│   350 │   │   │   fn_kwargs = self._prepare_subcommand_kwargs(self.subcommand)                   │
│ ❱ 351 │   │   │   fn(**fn_kwargs)                                                                │
│   352 │   │   else:                                                                              │
│   353 │   │   │   self.config_init = self.parser.instantiate_classes(self.config)                │
│   354 │   │   │   getattr(self, f"{self.subcommand}")()                                          │
│                                                                                                  │
│ /home/ashwin/projects/anomalib_copy/src/anomalib/engine/engine.py:923 in export                  │
│                                                                                                  │
│   920 │   │                                                                                      │
│   921 │   │   exported_model_path: Path | None = None                                            │
│   922 │   │   if export_type == ExportType.TORCH:                                                │
│ ❱ 923 │   │   │   exported_model_path = export_to_torch(                                         │
│   924 │   │   │   │   model=model,                                                               │
│   925 │   │   │   │   export_root=export_root,                                                   │
│   926 │   │   │   │   transform=transform,                                                       │
│                                                                                                  │
│ /home/ashwin/projects/anomalib_copy/src/anomalib/deploy/export.py:176 in export_to_torch         │
│                                                                                                  │
│   173 │   export_root = _create_export_root(export_root, ExportType.TORCH)                       │
│   174 │   metadata = get_metadata(task=task, model=model)                                        │
│   175 │   pt_model_path = export_root / "model.pt"                                               │
│ ❱ 176 │   torch.save(                                                                            │
│   177 │   │   obj={"model": inference_model, "metadata": metadata},                              │
│   178 │   │   f=pt_model_path,                                                                   │
│   179 │   )                                                                                      │
│                                                                                                  │
│ /home/ashwin/miniconda3/envs/temp/lib/python3.10/site-packages/torch/serialization.py:619 in     │
│ save                                                                                             │
│                                                                                                  │
│    616 │                                                                                         │
│    617 │   if _use_new_zipfile_serialization:                                                    │
│    618 │   │   with _open_zipfile_writer(f) as opened_zipfile:                                   │
│ ❱  619 │   │   │   _save(obj, opened_zipfile, pickle_module, pickle_protocol, _disable_byteorde  │
│    620 │   │   │   return                                                                        │
│    621 │   else:                                                                                 │
│    622 │   │   with _open_file_like(f, 'wb') as opened_file:                                     │
│                                                                                                  │
│ /home/ashwin/miniconda3/envs/temp/lib/python3.10/site-packages/torch/serialization.py:831 in     │
│ _save                                                                                            │
│                                                                                                  │
│    828 │   data_buf = io.BytesIO()                                                               │
│    829 │   pickler = pickle_module.Pickler(data_buf, protocol=pickle_protocol)                   │
│    830 │   pickler.persistent_id = persistent_id                                                 │
│ ❱  831 │   pickler.dump(obj)                                                                     │
│    832 │   data_value = data_buf.getvalue()                                                      │
│    833 │   zip_file.write_record('data.pkl', data_value, len(data_value))                        │
│    834                                                                                           │
│                                                                                                  │
│ /home/ashwin/miniconda3/envs/temp/lib/python3.10/site-packages/torch/serialization.py:786 in     │
│ persistent_id                                                                                    │
│                                                                                                  │
│    783 │   │   # see                                                                             │
│    784 │   │   # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external  │
│    785 │   │   # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537           │
│ ❱  786 │   │   if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj):          │
│    787 │   │   │                                                                                 │
│    788 │   │   │   if isinstance(obj, torch.storage.TypedStorage):                               │
│    789 │   │   │   │   # TODO: Once we decide to break serialization FC, this case               │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
RecursionError: maximum recursion depth exceeded while calling a Python object

Code of Conduct

• I agree to follow this project's Code of Conduct

[samet-akcay]

@adrianboguszewski any idea about why this is happening ?

[adrianboguszewski]

I saw this error a few times when exporting Torch models to OV, but I don't know the root cause or solution. I guess it's related to TorchScript or TorchDynamo.