v1.1.1

jmschrei · Oct 18, 2024 · 8e1a71e · 8e1a71e
1 parent 34a6e9a
commit 8e1a71e
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 2 deletions.
diff --git a/docs/whats_new.rst b/docs/whats_new.rst
@@ -5,6 +5,16 @@
 Release History
 ===============
 
+Version 1.1.1
+==============
+
+
+Highlights
+----------
+
+	- Fixed an issue with categorical distributions being used on the GPU
+
+
 Version 1.0.4
 ==============
 

diff --git a/pomegranate/distributions/categorical.py b/pomegranate/distributions/categorical.py
@@ -182,7 +182,9 @@ def log_probability(self, X):
 			max_value=self.n_keys-1, ndim=2, shape=(-1, self.d),
 			check_parameter=self.check_data)
 
-		logps = torch.zeros(X.shape[0], dtype=self.probs.dtype)
+		logps = torch.zeros(X.shape[0], dtype=self.probs.dtype, 
+			device=self.device)
+
 		for i in range(self.d):
 			if isinstance(X, torch.masked.MaskedTensor):
 				logp_ = self._log_probs[i][X[:, i]._masked_data]

diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
 	name='pomegranate',
-	version='1.1.0',
+	version='1.1.1',
 	author='Jacob Schreiber',
 	author_email='[email protected]',
 	packages=['pomegranate', 'pomegranate.distributions', 'pomegranate.hmm'],