PyTorch

import torch
import sol
import torchvision.models as models

''' Optimizing Model '''
py_model  = models.__dict__["alexnet"]()
input	  = torch.rand(32, 3, 224, 224)

sol_model = sol.optimize(py_model)

''' Run training '''
sol_model.train()

# You cannot initialize the optimizer at this point. You need to wait until
# you have executed the model at least once, so SOL has compiled it.
optimizer = None
for batch in ...:
	input, target = ...
	output = sol_model(input)
	loss = loss_func(output, target)
	# After running the model once, you can safely initialize the optimizer
	if optimizer is None:
		optimizer = torch.optim.Adam(sol_model.parameters(), ...)
	optimizer.zero_grad()
	loss.backward()
	optimizer.step()
	...

''' Run validation '''
sol_model.eval()
with torch.no_grad():
	for batch in ...:
		input = ...
		output = sol_model(input)
		...

model.predict = torch.compile(model.predict, backend='sol')

sol_model = sol.optimize(pytorch_model, mode='reduce-overhead')
# or
sol_model = torch.compile(pytorch_model, mode='reduce-overhead', backend='sol')

# Storing
sol_model = sol.optimize(pytorch_model, [...])
torch.save(sol_model.state_dict(), PATH)

# Loading
sol_model = sol.optimize(pytorch_model)
sol_model.load_state_dict(torch.load(PATH))

from torchvision import models

class Wrap(torch.nn.Module):
	def __init__(self, model):
		super().__init__()
		self.model = model

	def forward(self, x):
		out = self.model(x)
		if torch.jit.is_scripting():
			return (out[0], out[1]) if self.training else (out[0], None)
		return (out[0], out[1]) if self.training else (out, None)

model = Wrap(models.inception_v3())

# use only one output
model.training = False
sol_model = sol.optimize(model, ...)

# use two outputs
model.training = True
sol_model = sol.optimize(model, ...)

class ResNet50(L.LightningModule):
	def __init__(self):
		super().__init__()
		self.model = sol.optimize(torchvision.models.resnet50())

	def training_step(self, batch, batch_idx):
		...
	
	def configure_optimizers(self):
		...