Merge pull request #5 from braun-steven/fix/categorical-suff-stats-dim-3

Fix missing num_var parameter in suff_stat reshape

src/EinsumNetwork/ExponentialFamilyArray.py

@@ -532,7 +532,7 @@ def sufficient_statistics(self, x):
         if len(x.shape) == 2:
             stats = one_hot(x.long(), self.K)
         elif len(x.shape) == 3:
-            stats = one_hot(x.long(), self.K).reshape(-1, self.num_dims * self.K)
+            stats = one_hot(x.long(), self.K).reshape(-1, self.num_var, self.num_dims * self.K)
         else:
             raise AssertionError("Input must be 2 or 3 dimensional tensor.")
         return stats