Here is the modification I made, and it works.
Since prox elements are only seen as functions in the end, there is no need to put them in a ModuleList. But you can still do it.
class f(torch.nn.Module) :
    def __init__(self, prox_el) :
        super().__init__()
        self.prox_el = prox_el

    def forward(self, x, s=None) :
        return self.prox_el(x.unsqueeze(0)).squeeze()

class PGD(torch.optim.Optimizer):
    """Proximal Gradient Descent
    Args:
      params: [torch.Parameter]
        List of parameters to optimize over
      prox: [callable or None]
        List of prox operators, one per parameter.
      lr: float
        Learning rate
      momentum: float in [0, 1]
      normalization: str
        Type of gradient normalization to be used.
        Possible values are 'none', 'L2', 'Linf', 'sign'.
    """
    name = 'PGD'
    POSSIBLE_NORMALIZATIONS = {'none', 'L2', 'Linf', 'sign'}

    def __init__(self, params, prox=None, lr=.1, momentum=.9, normalization='none'):
        if prox is None:
            prox = [None] * len(list(params))

        self.prox = []
        for prox_el in prox:
            if prox_el is not None:
                self.prox.append(f(prox_el))
            else:
                self.prox.append(torch.nn.Identity())
        #...
        # ...