Unable to script model

[callmekris]

trafficstars

Currently running mosaic composer in a databricks environment but experiencing different behavior of torch.jit.script after importing composer depending on the DBR installed.

Code to reproduce

import torch
class MyModule(torch.nn.Module):
    def forward(self, x):
        return x + 10

m = torch.jit.script(MyModule())
print(m)
print(type(m))

and

import torch
from composer.models import ComposerModel
class MyModule(torch.nn.Module):
    def forward(self, x):
        return x + 10

m = torch.jit.script(MyModule())
print(m)
print(type(m))

This is just a minimum example to reproduce, but the same happens when training a composer model and using export_for_inference.

Expected Behavior DBR 14.3

Running with both with or without the composer import prints as expected:

RecursiveScriptModule(original_name=MyModule)
<class 'torch.jit._script.RecursiveScriptModule'>

On DBR >= 15.1

Running with the composer import prints:

MyModule()
<class '__main__.MyModule'>

And without the composer import it scripts correctly as expected.

On both runtimes the libary versions are: torch 2.4.0 mosaicml 0.24.1 mosaicml-cli 0.6.41

I'm not exactly sure how to pinpoint the cause it only occurs on the newer DBRs even 15.4 LTS, but only after the import of composer.