def _process_stochastic(output: torch.Tensor, parents: [storch.Tensor],
plates: [storch.Plate]):
if isinstance(output, storch.Tensor):
if not output.stochastic:
# TODO: Calls _add_parents so something is going wrong here
# The Tensor was created by calling @deterministic within a stochastic context.
# This means that we have to conservatively assume it is dependent on the parents
output._add_parents(storch.wrappers._stochastic_parents)
return output
if isinstance(output, torch.Tensor):
t = storch.Tensor(output, parents, plates)
return t
else:
raise TypeError(
"All outputs of functions wrapped in @storch.stochastic "
"should be Tensors. At " + str(output))
