Integrate comment of passive aggressive reviewer

aleximmer · Dec 10, 2021 · 5ee4009 · 5ee4009
1 parent d37f26b
commit 5ee4009
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/laplace/baselaplace.py b/laplace/baselaplace.py
@@ -885,7 +885,11 @@ def fit(self, train_loader, override=True):
 
         X, _ = next(iter(train_loader))
         with torch.no_grad():
-            self.n_outputs = self.model(X[:1].to(self._device)).shape[-1]
+            try:
+                out = self.model(X[:1].to(self._device))
+            except (TypeError, AttributeError):
+                out = self.model(X.to(self._device))
+        self.n_outputs = out.shape[-1]
         setattr(self.model, 'output_size', self.n_outputs)
 
         eigenvectors, eigenvalues, loss = self.backend.eig_lowrank(train_loader)