fix: small fix to expectation layer

qiboteam · Jul 17, 2024 · 42c0a1b · 42c0a1b
1 parent 2e0c1bd
commit 42c0a1b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/src/qiboml/models/encoding_decoding.py b/src/qiboml/models/encoding_decoding.py
@@ -111,9 +111,10 @@ def __post_init__(self):
 
     def forward(self, x: Circuit) -> "ndarray":
         if self.nshots is None:
+            # breakpoint()
             return self.observable.expectation(
                 super().forward(x).state(),
-            )
+            ).reshape(1, 1)
         else:
             return self.observable.expectation_from_samples(
                 super().forward(x).samples(),

diff --git a/tests/test_models_pytorch.py b/tests/test_models_pytorch.py
@@ -17,7 +17,7 @@ def get_layers(module, layer_type=None):
     for name, layer in inspect.getmembers(module, inspect.isclass):
         if layer.__module__ == module.__name__:
             if layer_type is not None:
-                if issubclass(layer, layer_type):
+                if issubclass(layer, layer_type) and not layer is layer_type:
                     pass
                 else:
                     continue