TorchScript bad_alloc issue (#2542)

Summary: Pull Request resolved: #2542

Differential Revision: D65495806

Author: PaulZhang12

torchrec/models/tests/test_deepfm.py

@@ -202,7 +202,7 @@ def test_fx_script(self) -> None:
             sparse_features=sparse_features,
         )
 
-        gm = symbolic_trace(deepfm_nn)
+        gm = torch.fx.GraphModule(deepfm_nn, Tracer().trace(deepfm_nn))
 
         scripted_gm = torch.jit.script(gm)
 

torchrec/modules/tests/test_mc_modules.py

@@ -11,6 +11,7 @@
 from typing import Dict
 
 import torch
+from torchrec.fx import Tracer
 from torchrec.modules.mc_modules import (
     average_threshold_filter,
     DistanceLFU_EvictionPolicy,
@@ -357,5 +358,5 @@ def test_fx_jit_script_not_training(self) -> None:
         )
 
         model.train(False)
-        gm = torch.fx.symbolic_trace(model)
+        gm = torch.fx.GraphModule(model, Tracer().trace(model))
         torch.jit.script(gm)

torchrec/modules/tests/test_mlp.py

@@ -14,7 +14,7 @@
 import torch
 from hypothesis import given, settings
 from torch import nn
-from torchrec.fx import symbolic_trace
+from torchrec.fx import symbolic_trace, Tracer
 from torchrec.modules.mlp import MLP, Perceptron
 
 
@@ -99,13 +99,13 @@ def test_fx_script_Perceptron(self) -> None:
         # Dry-run to initialize lazy module.
         m(torch.randn(batch_size, in_features))
 
-        gm = symbolic_trace(m)
+        gm = torch.fx.GraphModule(m, Tracer().trace(m))
         torch.jit.script(gm)
 
     def test_fx_script_MLP(self) -> None:
         in_features = 3
         layer_sizes = [16, 8, 4]
         m = MLP(in_features, layer_sizes)
 
-        gm = symbolic_trace(m)
+        gm = torch.fx.GraphModule(m, Tracer().trace(m))
         torch.jit.script(gm)