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!