# Tensor Sizes

When writing your model in PyTorch, be aware that when this model is run inside Nuke, the input tensor to your model forward function will be of size **1 x inChan x inH x inW**, where **inH** and **inW** are the height and width of the image passed to the Inference node.

Similarly, the output tensor that is returned from your model forward function must be of size **1 x outChan x outH x outW**, where **outH** and **outW** are the height and width of the image output from the Inference node. **inChan** and **outChan** are the number of channels you have defined in the **Channels In** and **Channels Out** knobs in the CatFileCreator respectively.

For example, in our simple Cifar10ResNet example, the Cifar10ResNet model processed an RGB image and returned a single channel output image of ones/zeros. We set up our CatFileCreator node as follows:

When this model is run inside Nuke, the shape of the **input** tensor passed to the forward function:

```
def forward(self, input):
"""
:param input: A torch.Tensor of size 1 x 3 x H x W representing the input image
:return: A torch.Tensor of size 1 x 1 x H x W of zeros or ones
"""
modelOutput = self.model.forward(input)
modelLabel = int(torch.argmax(modelOutput[0]))
plane = 0
if modelLabel == plane:
output = torch.ones(1, 1, input.shape[2], input.shape[3])
else:
output = torch.zeros(1, 1, input.shape[2], input.shape[3])
return output
```

will be of size **1 x 3 x H x W** and the **output** tensor returned from the forward function must be of size **1 x 1 x H x W**.

If we’re using a model that processes 5 channels and outputs 3 channels, and whose **Output Scale** value is 2:

When this model is run inside Nuke, the **input** tensor to the model forward function will be of size **1 x 5 x inH x inW** and its **output** tensor must be of size **1 x 3 x 2*inH x 2*inW**.