remove mapping reduced from tests, hoping we wont need to register bu…

…ffer
FAIR-Chem · Sep 20, 2024 · dcd2501 · dcd2501
1 parent 4ff238c
commit dcd2501
Showing 1 changed file with 0 additions and 7 deletions.
diff --git a/tests/core/models/test_escn_compiles.py b/tests/core/models/test_escn_compiles.py
@@ -26,7 +26,6 @@
 from fairchem.core.datasets import data_list_collater
 from fairchem.core.models.escn import escn_exportable
 from fairchem.core.models.escn.so3_exportable import (
-    CoefficientMapping,
     SO3_Grid,
 )
 from fairchem.core.models.scn.smearing import GaussianSmearing
@@ -210,7 +209,6 @@ def test_escn_so2_conv_exports_and_compiles(self, tol=1e-5) -> None:
         }
 
         lmax, mmax = 4, 2
-        mappingReduced = escn_exportable.CoefficientMapping([lmax], [mmax])
         shpere_channels = 128
         edge_channels = 128
         args = (torch.rand(680, 19, shpere_channels), torch.rand(680, edge_channels))
@@ -222,7 +220,6 @@ def test_escn_so2_conv_exports_and_compiles(self, tol=1e-5) -> None:
             lmax=lmax,
             mmax=mmax,
             act=torch.nn.SiLU(),
-            mappingReduced=mappingReduced,
         )
         prog = export(so2, args=args, dynamic_shapes=dynamic_shapes1)
         export_out = prog.module()(*args)
@@ -244,7 +241,6 @@ def test_escn_message_block_exports_and_compiles(self, tol=1e-5) -> None:
         SO3_grid = torch.nn.ModuleDict()
         SO3_grid["lmax_lmax"] = SO3_Grid(lmax, lmax)
         SO3_grid["lmax_mmax"] = SO3_Grid(lmax, mmax)
-        mappingReduced = CoefficientMapping([lmax], [mmax])
         message_block = escn_exportable.MessageBlock(
             layer_idx=0,
             sphere_channels=sphere_channels,
@@ -256,7 +252,6 @@ def test_escn_message_block_exports_and_compiles(self, tol=1e-5) -> None:
             max_num_elements=90,
             SO3_grid=SO3_grid,
             act=torch.nn.SiLU(),
-            mappingReduced=mappingReduced,
         )
 
         # generate inputs
@@ -297,7 +292,6 @@ def test_escn_layer_block_exports_and_compiles(self, tol=1e-5) -> None:
         SO3_grid = torch.nn.ModuleDict()
         SO3_grid["lmax_lmax"] = SO3_Grid(lmax, lmax)
         SO3_grid["lmax_mmax"] = SO3_Grid(lmax, mmax)
-        mappingReduced = CoefficientMapping([lmax], [mmax])
         layer_block = escn_exportable.LayerBlock(
             layer_idx=0,
             sphere_channels=sphere_channels,
@@ -309,7 +303,6 @@ def test_escn_layer_block_exports_and_compiles(self, tol=1e-5) -> None:
             max_num_elements=90,
             SO3_grid=SO3_grid,
             act=torch.nn.SiLU(),
-            mappingReduced=mappingReduced,
         )
 
         # generate inputs