Advice on how to use a self-supervised regression scheme within a single step in pl
Hi
I have the following scheme:
class refine_P(LightningModule):

    def __init__(
        self, hparams,
    ):
        self.model = #
        self.regressor #
   def training_step(self, batch, batch_idx, is_train=True):
      out = self.model(batch)
      self.regressor.reset_parameters()
      base_loss = self.base_loss(out)
      while condition:
         out = self.regressor(out)
         loss = self.regress_loss(out)
         loss.backward()
         ...
      return base_loss

I.e - Without using pl, it will be just two training procedures, one within the other.
My question is, Can I somehow create a trainer within this module, and call this trainer each training step of refine_P?