Unable to script model #3584
Description
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.