Repository: Lightning-Universe/lightning-flash

Issue:
Image classification `NoModule` bug with multi-gpu
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

Run flash image classification with multiple GPUs, bug in NoModule getattr
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->


#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Environment

 - PyTorch Version (e.g., 1.0):
 - OS (e.g., Linux):
 - How you installed PyTorch (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version:
 - GPU models and configuration:
 - Any other relevant information:

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:

Resolving patch:
diff --git a/flash/image/classification/adapters.py b/flash/image/classification/adapters.py
--- a/flash/image/classification/adapters.py
+++ b/flash/image/classification/adapters.py
@@ -57,25 +57,6 @@ def remap(self, data, mapping):
         return data
 
 
-class NoModule:
-
-    """This class is used to prevent nn.Module infinite recursion."""
-
-    def __init__(self, task):
-        self.task = task
-
-    def __getattr__(self, key):
-        if key != "task":
-            return getattr(self.task, key)
-        return self.task
-
-    def __setattr__(self, key: str, value: Any) -> None:
-        if key == "task":
-            object.__setattr__(self, key, value)
-            return
-        setattr(self.task, key, value)
-
-
 class Model(torch.nn.Module):
     def __init__(self, backbone: torch.nn.Module, head: Optional[torch.nn.Module]):
         super().__init__()
@@ -97,7 +78,6 @@ class Learn2LearnAdapter(Adapter):
 
     def __init__(
         self,
-        task: AdapterTask,
         backbone: torch.nn.Module,
         head: torch.nn.Module,
         algorithm_cls: Type[LightningModule],
@@ -143,7 +123,6 @@ def __init__(
 
         super().__init__()
 
-        self._task = NoModule(task)
         self.backbone = backbone
         self.head = head
         self.algorithm_cls = algorithm_cls
@@ -309,7 +288,9 @@ def from_task(
                 "The `shots` should be provided training_strategy_kwargs={'shots'=...}. "
                 "This is equivalent to the number of sample per label to select within a task."
             )
-        return cls(task, backbone, head, algorithm, **kwargs)
+        adapter = cls(backbone, head, algorithm, **kwargs)
+        adapter.__dict__["_task"] = task
+        return adapter
 
     def training_step(self, batch, batch_idx) -> Any:
         input = (batch[DataKeys.INPUT], batch[DataKeys.TARGET])
@@ -337,7 +318,7 @@ def _sanetize_batch_size(self, batch_size: int) -> int:
             warning_cache.warn(
                 "When using a meta-learning training_strategy, the batch_size should be set to 1. "
                 "HINT: You can modify the `meta_batch_size` to 100 for example by doing "
-                f"{type(self._task.task)}" + "(training_strategies_kwargs={'meta_batch_size': 100})"
+                f"{type(self._task)}" + "(training_strategies_kwargs={'meta_batch_size': 100})"
             )
         return 1
 
@@ -486,10 +467,9 @@ class DefaultAdapter(Adapter):
 
     required_extras: str = "image"
 
-    def __init__(self, task: AdapterTask, backbone: torch.nn.Module, head: torch.nn.Module):
+    def __init__(self, backbone: torch.nn.Module, head: torch.nn.Module):
         super().__init__()
 
-        self._task = NoModule(task)
         self.backbone = backbone
         self.head = head
 
@@ -503,23 +483,25 @@ def from_task(
         head: torch.nn.Module,
         **kwargs,
     ) -> Adapter:
-        return cls(task, backbone, head)
+        adapter = cls(backbone, head)
+        adapter.__dict__["_task"] = task
+        return adapter
 
     def training_step(self, batch: Any, batch_idx: int) -> Any:
         batch = (batch[DataKeys.INPUT], batch[DataKeys.TARGET])
-        return Task.training_step(self._task.task, batch, batch_idx)
+        return Task.training_step(self._task, batch, batch_idx)
 
     def validation_step(self, batch: Any, batch_idx: int) -> Any:
         batch = (batch[DataKeys.INPUT], batch[DataKeys.TARGET])
-        return Task.validation_step(self._task.task, batch, batch_idx)
+        return Task.validation_step(self._task, batch, batch_idx)
 
     def test_step(self, batch: Any, batch_idx: int) -> Any:
         batch = (batch[DataKeys.INPUT], batch[DataKeys.TARGET])
-        return Task.test_step(self._task.task, batch, batch_idx)
+        return Task.test_step(self._task, batch, batch_idx)
 
     def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
         batch[DataKeys.PREDS] = Task.predict_step(
-            self._task.task, (batch[DataKeys.INPUT]), batch_idx, dataloader_idx=dataloader_idx
+            self._task, (batch[DataKeys.INPUT]), batch_idx, dataloader_idx=dataloader_idx
         )
         return batch
 
diff --git a/flash/pointcloud/detection/open3d_ml/__init__.py b/flash/pointcloud/detection/open3d_ml/__init__.py
new file mode 100644


Test patch: