Custom Optimizers with lightning

1

I am trying to use this custom optimizer: GitHub - ItzikMalkiel/MTAdam: MTAdam: Automatic Balancing of Multiple Training Loss Terms

However I can’t seem to figure out how to integrate it with lighting. I tried using Manual Optimization, but could not get it to work. The challenges are:

  1. It has a custom step function that requires the loss be passed in
  2. It calls backwards inside step

Any idea how I could use this with lightning? Or am I better off in this case sticking to regular pytorch?