device none if layer is not module

microsoft · Nov 12, 2024 · 0ec0b1e · 0ec0b1e
1 parent d3d2096
commit 0ec0b1e
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mttl/models/containers/base.py b/mttl/models/containers/base.py
@@ -54,7 +54,10 @@ def __init__(self, config, layer, selector=None):
         self._enabled = True
         self.config = config
         self.layer = layer
-        self.device = next(layer.parameters()).device
+        if isinstance(layer, nn.Module):
+            self.device = next(layer.parameters()).device
+        else:
+            self.device = None
         self.selector = selector or TaskNameSelector()
         self._default_expert_name = None
         self.expert_infos = {}