def set_output(self, tensor, buffer): if isinstance(tensor, edsl.Tensor): tensor = edsl.TensorRef(tensor) self.outputs[tensor] = buffer return self
def output(self, tensor): if isinstance(tensor, edsl.Tensor): tensor = edsl.TensorRef(tensor) return self.outputs.get(tensor)