Casted sample_weights as tensor as well (#1104)

jmschrei · Jul 11, 2024 · badd707 · badd707
1 parent 9a8cf3b
commit badd707
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pomegranate/distributions/normal.py b/pomegranate/distributions/normal.py
@@ -257,7 +257,7 @@ def summarize(self, X, sample_weight=None):
 
 		X, sample_weight = super().summarize(X, sample_weight=sample_weight)
 		X = _cast_as_tensor(X, dtype=self.means.dtype)
-
+		sample_weight = _cast_as_tensor(sample_weight, dtype=self.means.dtype)
 		if self.covariance_type == 'full':
 			self._w_sum += torch.sum(sample_weight, dim=0)
 			self._xw_sum += torch.sum(X * sample_weight, axis=0)