My current solution is modifying the class Foo as follows:
class Foo:
    def __init__(self, x: Tensor):
        self.x = x
    def transform(self, p: Tensor):
        self.x = self.x.detach()
        self.x = self.x + p
    def loss(self):
        return self.x.norm()**2