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)