Update compile test

Author: qubvel

docs/encoders.rst

@@ -97,14 +97,14 @@ densenet201                  imagenet                             18M         
 densenet161                  imagenet                             26M         ✅        ✅         ✅       
 inceptionresnetv2            imagenet / imagenet+background       54M         ✅        ✅         ✅       
 inceptionv4                  imagenet / imagenet+background       41M         ✅        ✅         ✅       
-efficientnet-b0              imagenet / advprop                   4M          ❌        ❌         ✅       
-efficientnet-b1              imagenet / advprop                   6M          ❌        ❌         ✅       
-efficientnet-b2              imagenet / advprop                   7M          ❌        ❌         ✅       
-efficientnet-b3              imagenet / advprop                   10M         ❌        ❌         ✅       
-efficientnet-b4              imagenet / advprop                   17M         ❌        ❌         ✅       
-efficientnet-b5              imagenet / advprop                   28M         ❌        ❌         ✅       
-efficientnet-b6              imagenet / advprop                   40M         ❌        ❌         ✅       
-efficientnet-b7              imagenet / advprop                   63M         ❌        ❌         ✅       
+efficientnet-b0              imagenet / advprop                   4M          ❌        ✅         ✅       
+efficientnet-b1              imagenet / advprop                   6M          ❌        ✅         ✅       
+efficientnet-b2              imagenet / advprop                   7M          ❌        ✅         ✅       
+efficientnet-b3              imagenet / advprop                   10M         ❌        ✅         ✅       
+efficientnet-b4              imagenet / advprop                   17M         ❌        ✅         ✅       
+efficientnet-b5              imagenet / advprop                   28M         ❌        ✅         ✅       
+efficientnet-b6              imagenet / advprop                   40M         ❌        ✅         ✅       
+efficientnet-b7              imagenet / advprop                   63M         ❌        ✅         ✅       
 mobilenet_v2                 imagenet                             2M          ✅        ✅         ✅       
 xception                     imagenet                             20M         ✅        ✅         ✅       
 timm-efficientnet-b0         imagenet / advprop / noisy-student   4M          ✅        ✅         ✅       

segmentation_models_pytorch/encoders/efficientnet.py

@@ -35,9 +35,6 @@
 class EfficientNetEncoder(EfficientNet, EncoderMixin):
     _is_torch_scriptable = False
 
-    # works with torch 2.4.0, but not with torch 2.5.1
-    _is_torch_compilable = False
-
     def __init__(
         self,
         stage_idxs: List[int],

tests/encoders/base.py

@@ -212,6 +212,7 @@ def test_compile(self):
         encoder = self.get_tiny_encoder()
         encoder = encoder.eval().to(default_device)
 
+        torch.compiler.reset()
         compiled_encoder = torch.compile(encoder, fullgraph=True, dynamic=True)
 
         if encoder._is_torch_compilable:

tests/models/base.py

@@ -228,8 +228,10 @@ def test_compile(self):
             self.skipTest("No diff and not on `main`.")
 
         sample = self._get_sample().to(default_device)
-
         model = self.get_default_model()
+        model = model.eval().to(default_device)
+
+        torch.compiler.reset()
         compiled_model = torch.compile(model, fullgraph=True, dynamic=True)
 
         with torch.inference_mode():