example_model.py

model = resnet50(pretrained=True)
# change input channels number to match the rasterizer's output
model.conv1 = nn.Conv2d(
   num_in_channels,
   model.conv1.out_channels,
   kernel_size=model.conv1.kernel_size,
   stride=model.conv1.stride,
   padding=model.conv1.padding,
   bias=False,
)
# change output size to (X, Y) * number of future states
model.fc = nn.Linear(in_features=2048, out_features=num_targets)