Torchscript compatibility with control flows

LewsTherin: <p>I’m writing a deep learning model in pytorch…in the constructor, I’m defining a transformation function (the <code>self.transform</code>) which will be applied to the input before the forward pass.</p>
<pre><code class="lang-auto">class CustomModel(torch.nn.Module):

    def __init__(self, model, model_name: str, model_params: dict):
        super().__init__()
        self.model_name = model_name
        self.model_params = model_params
        self.model = model

        # Generate the sequential transformations needed for each clip
        self.transform = torch.nn.Sequential(
            LongSideScale(size=self.model_params["side_size"]),
            CenterCrop(size=(self.model_params["crop_size"], self.model_params["crop_size"])),
            Div255(),
            Normalize(self.model_params["mean"], self.model_params["std"]),
        )

    def forward(self, batch_video: List[torch.Tensor]) -&gt; torch.Tensor:
        batch_video = self.apply_batch_transform(batch_video)
        logits = self.model(batch_video)
        probability = torch.sigmoid(logits)
        return probability

    def apply_batch_transform(self, batch_video: List[torch.Tensor]) -&gt; torch.Tensor:
        transformed_batch_video = [self.transform(clip) for clip in batch_video]
        transformed_batch_video = torch.stack(transformed_batch_video)
        return transformed_batch_video
</code></pre>
<p>The problem is that the input can be either:</p>
<ul>
<li>sequence of N RGB frames: x.shape = (3, N, w, h)</li>
<li>frames can be stacked: x.shape = (3*N, w, h)</li>
</ul>
<p>This is accounted for in the implementation of the <code>LongSideScale</code> part of the transformation:</p>
<pre><code class="lang-auto">def long_side_scale(
    x: torch.Tensor,
    size: int,
    interpolation: str = "bilinear",
    backend: str = "pytorch",
) -&gt; torch.Tensor:
    """
    Determines the longer spatial dim of the video (i.e. width or height) and scales
    it to the given size. To maintain aspect ratio, the shorter side is then scaled
    accordingly.
    """

    # Check if tensor is 3D (C, H, W)
    if len(x.shape) == 3:
        c, h, w = x.shape
        x = x.unsqueeze(1)  # Add a dummy temporal dimension (T=1)
        unpack = True
    elif len(x.shape) == 4:
        c, t, h, w = x.shape
        unpack = False
    else:
        raise ValueError(f"Input tensor must be 3D or 4D, but got shape {x.shape}")

    is_uint8 = x.dtype == torch.uint8
    c, t, h, w = x.shape
    if w &gt; h:
        new_h = int(math.floor((float(h) / w) * size))
        new_w = size
    else:
        new_h = size
        new_w = int(math.floor((float(w) / h) * size))

    if backend == "pytorch":
        if is_uint8:
            x = x.to(torch.float32)
        x = torch.nn.functional.interpolate(x, size=(new_h, new_w), mode=interpolation, align_corners=False)
        if is_uint8:
            x = x.to(torch.uint8)

        # If we added a temporal dimension, squeeze it back
        if unpack:
            x = x.squeeze(1)
        return x
    else:
        raise NotImplementedError(f"{backend} backend not supported.")
</code></pre>
<p>Finally, this model will be converted to torchscript, so something like:</p>
<pre><code class="lang-auto">model = CustomModel(model=model, model_name=model_name, model_params=model_params)
scripted = torch.jit.script(model)
scripted.save(scripted_path)
</code></pre>
<p>I was worried this could not be the best solution, since TorchScript could have limitations regarding dynamic control flow, such as if conditions based on tensor properties or data types, when converting to a static graph (I could be wrong, though!).<br>
I tested the code and everything seems to be working fine, the model is saved an I can perform inference for both input scenarios. I was just wondering if it ight be an issue, and if there’s a better way to tackle this situation.</p>