Fix EmpiricalMarginal with multiple GPUs
@@ -51,8 +51,11 @@ class EmpiricalMarginal(Empirical):
samples_by_chain = []
weights_by_chain = []
for i in range(num_chains):
- samples_by_chain.append(torch.stack(self._samples_buffer[i], dim=0))
- weights_by_chain.append(torch.stack(self._weights_buffer[i], dim=0))
+ samples = torch.stack(self._samples_buffer[i], dim=0)
+ samples_by_chain.append(samples)
+ weights_dtype = samples.dtype if samples.dtype.is_floating_point else torch.float32
+ weights = torch.as_tensor(self._weights_buffer[i], device=samples.device, dtype=weights_dtype)
+ weights_by_chain.append(weights)
if len(samples_by_chain) == 1:
return samples_by_chain[0], weights_by_chain[0]
else:
@@ -73,14 +76,10 @@ class EmpiricalMarginal(Empirical):
in ``[0, num_chains - 1]``, and there must be equal number
of samples per chain.
"""
- weight_type = value.new_empty(1).float().type() if value.dtype in (torch.int32, torch.int64) \
- else value.type()
# Apply default weight of 1.0.
if log_weight is None:
- log_weight = torch.tensor(0.0).type(weight_type)
- if isinstance(log_weight, numbers.Number):
- log_weight = torch.tensor(log_weight).type(weight_type)
- if self._validate_args and log_weight.dim() > 0:
+ log_weight = 0.0
+ if self._validate_args and not isinstance(log_weight, numbers.Number) and log_weight.dim() > 0:
raise ValueError("``weight.dim() > 0``, but weight should be a scalar.")
# Append to the buffer list
