torch.jit.load causes an error when using scalene

[chaitanya-kolluru-elucid]

Describe the bug
I have some code that loads a pytorch model with torch.jit.load. Things run fine when running the script normally, but when adding scalene with python -m scalene <program.py> it breaks.

Screenshots

  File "/usr/local/lib/python3.8/site-packages/torch/jit/_serialization.py", line 169, in load
    return wrap_cpp_module(cpp_module)
  File "/usr/local/lib/python3.8/site-packages/torch/jit/_recursive.py", line 885, in wrap_cpp_module
    return torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
AttributeError: type object 'RecursiveScriptModule' has no attribute '_construct'

[emeryberger]

We need a minimum working example so we can reproduce and debug this issue.

[Issikul]

Try downgrading to an earlier version. Version 1.5.51 gives me the same error message when loading a torch.jit model, even though its not in the scope of the profiling.

[emeryberger]

Can you provide a small code snippet that demonstrates the issue?

[Issikul]

Sure! Here is the majorly simplifed code I'm trying to run:

import torch


class PythonModel:

    def initialize(self):
        self.device = "cuda"
        self._load_models(
            decoder_model_path="/path/to/torchscipt/model.ts",
            device=self.device
        )

    def execute(self, data):
        results = []

        # DO STUFF

        return results

    def _load_models(self, model_path: str, device: str):
        self.decoder = torch.jit.optimize_for_inference(torch.jit.load(model_path, map_location=device))


if __name__ == "__main__":
    model = PythonModel()
    model.initialize()
    model.execute(data=None)

I'm using the following command to run the script (error trace at the end of my comment):

scalene model_test.py --html > my.html

(But its the same with scalene model_test.py)

Package versions (installed with pip):
torch 2.3.1
scalene 1.5.51

It's working fine with scalene-1.5.48.

Error in program being profiled:
 type object 'RecursiveScriptModule' has no attribute '_construct'
Traceback (most recent call last):
  File "/path/to/my/venv/site-packages/scalene/scalene_profiler.py", line 1767, in profile_code
    exec(code, the_globals, the_locals)
  File "/path/to/my/script/model_test.py", line 26, in <module>
    model.initialize()
  File "/path/to/my/script/model_test.py", line 8, in initialize
    self._load_models(
  File "/path/to/my/script/model_test.py", line 21, in _load_models
    self.decoder = torch.jit.optimize_for_inference(torch.jit.load(decoder_model_path, map_location=device))
  File "/path/to/my/venv/site-packages/torch/jit/_serialization.py", line 165, in load
    return wrap_cpp_module(cpp_module)
  File "/path/to/my/venv/site-packages/torch/jit/_recursive.py", line 1030, in wrap_cpp_module
    return torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
AttributeError: type object 'RecursiveScriptModule' has no attribute '_construct'

[emeryberger]

I need something I can run; this code, as is, does not. I had to modify the keyword on line 8 to model_path, but even with that modification, I need the model file. Can you upload a complete example in a zip file?

[sa1g]

Hi, I'm joining the thread.
I'm having the same error: it runs without scalene, some torch.jit functionalities disappear when profiling.

Here's a simple example:

import clip
import torch

if __name__ == "__main__":
    # Example usage
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, preprocess = clip.load("ViT-B/32", device=device, jit=False)
    print("Model loaded successfully.")

Requirements

Python 3.10.14 (tried 3.13, but clip is not available).

torch
clip
scalene

[Edit: I don't know if it somehow matters, this was tested using conda, everything installed with pip]

Scalene "commands"

I tried running with:

• python3 -m scalene <file_name>.py
• scalene file_name.py
• Importing and adding from scalene import scalene_profiler, scalene_profiler.<stop-start>() to ignore the model loading (I don't want to profile it, I need to profile the rest of the code) and starting with scalene --off <file_name>.py and the same with python's call.

Here's the error trace:

> python3 -m scalene temp.py 
Error in program being profiled:
 type object 'RecursiveScriptModule' has no attribute '_construct'
Traceback (most recent call last):
  File "/REDACTED/python3.10/site-packages/scalene/scalene_profiler.py", line 1767, in profile_code
    exec(code, the_globals, the_locals)
  File "/home/ettore/Projects/uni/dl/clip-tta/temp.py", line 7, in <module>
    model, preprocess = clip.load("ViT-B/32", device="cuda", jit=False)
  File "/REDACTED/python3.10/site-packages/clip/clip.py", line 129, in load
    model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
  File "/REDACTED/python3.10/site-packages/torch/jit/_serialization.py", line 165, in load
    return wrap_cpp_module(cpp_module)
  File "/REDACTED/python3.10/site-packages/torch/jit/_recursive.py", line 1030, in wrap_cpp_module
    return torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
AttributeError: type object 'RecursiveScriptModule' has no attribute '_construct'

[emeryberger]

Reproduced. Working on a solution.