Thanks for leaving the link to your notebook
I already had this error when defining one of the methods of my pytorch module as a lambda function, like self.proj = lambda x : x.
When I changed it to self.proj = nn.Identity() I didn’t get any error.
So the error could come from the fact that prox contains lambda functions.
You can change lambda x, s=None: x for example to an nn.Identity().
Try to remove the lambda functions, and use pytorch module.
Note : I get this when executing your code AttributeError: Can't pickle local object 'PGD.__init__.<locals>.<lambda>'