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))