PyTorch

import torch
import sol
import torchvision.models as models

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

# Use vdims=[True] if you plan to use changing batchsizes
sol_model = sol.optimize(py_model, vdims=[True])

''' Run training '''
sol_model.train()
for batch in ...:
	input, target = ...
	output = sol_model(input)
	loss = loss_func(output, target)
	loss.backward()
	...

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

# 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, ...)