In case you need to run layers that are not supported by SOL you can partially implement the model using SOL:
class MyModel(framework.Model): def __init__(self, ...): self.A = sol.optimize(PartThatCanBeOptimized, ...) self.B = PartThatCannotBeOptimized self.C = sol.optimize(OtherPartThatCanBeOptimized, ...) # don't forget to set the inputs of this submodel to requiresGrad=True! def forward(self, X): X = self.A(X) # executed by SOL X = self.B(X) # executed by Framework X = self.C(X) # executed by SOL return X
It is important that except the input part of the network is called with
sol.optimize(..., sol.input(..., requires_grad=True), ...) to ensure correct gradient calculations!