Module checks for torch.Distribution, makes better runtimeerrors
@@ -4,8 +4,8 @@ from collections import OrderedDict
import torch
from torch import nn
+from torch.distributions import Distribution
-from .distributions import Distribution
from .lazy import LazyTensor
from .variational import VariationalStrategy
@@ -87,12 +87,12 @@ class Module(nn.Module):
def named_hyperparameters(self):
for name, param in self.named_parameters():
- if 'variational_' not in name:
+ if "variational_" not in name:
yield name, param
def named_variational_parameters(self):
for name, param in self.named_parameters():
- if 'variational_' in name:
+ if "variational_" in name:
yield name, param
def named_variational_strategies(self):
@@ -178,18 +178,27 @@ class Module(nn.Module):
def __call__(self, *inputs, **kwargs):
outputs = self.forward(*inputs, **kwargs)
- if torch.is_tensor(outputs) or isinstance(outputs, Distribution) or isinstance(outputs, LazyTensor):
- return outputs
- for output in outputs:
- if not (isinstance(output, Distribution) or torch.is_tensor(output) or isinstance(output, LazyTensor)):
+
+ if isinstance(outputs, tuple):
+ if not all(
+ torch.is_tensor(output) or isinstance(output, Distribution) or isinstance(output, LazyTensor)
+ for output in outputs
+ ):
raise RuntimeError(
- "Output must be a Distribution, torch.Tensor, or LazyTensor. "
- "Was a {}".format(input.__class__.__name__)
+ "All outputs must be a Distribution, torch.Tensor, or LazyTensor. "
+ "Got {}".format([output.__class__.__name__ for output in outputs])
)
if len(outputs) == 1:
outputs = outputs[0]
return outputs
+ elif torch.is_tensor(outputs) or isinstance(outputs, Distribution) or isinstance(outputs, LazyTensor):
+ return outputs
+ else:
+ raise RuntimeError(
+ "Output must be a Distribution, torch.Tensor, or LazyTensor. Got {}".format(outputs.__class__.__name__)
+ )
+
def _extract_named_parameter_priors(module, memo=None, prefix=""):
if memo is None:
