huggingface/pytorch-image-models · error · Exception

Memory Efficient not supported in JIT

Error message

Memory Efficient not supported in JIT

What it means

DenseNet's gradient-checkpointing path cannot run under TorchScript: torch.jit.script of a module whose forward reaches the checkpoint branch raises Exception('Memory Efficient not supported in JIT'). The check fires when grad_checkpointing is enabled and any input requires grad while the model is being scripted/traced in JIT mode.

Source

Thrown at timm/models/densenet.py:100

    # torchscript does not yet support *args, so we overload method
    # allowing it to take either a List[Tensor] or single Tensor
    def forward(self, x: Union[torch.Tensor, List[torch.Tensor]]) -> torch.Tensor:  # noqa: F811
        """Forward pass.

        Args:
            x: Input features (single tensor or list of tensors).

        Returns:
            New features to be concatenated.
        """
        if isinstance(x, torch.Tensor):
            prev_features = [x]
        else:
            prev_features = x

        if self.grad_checkpointing and self.any_requires_grad(prev_features):
            if torch.jit.is_scripting():
                raise Exception("Memory Efficient not supported in JIT")
            bottleneck_output = self.call_checkpoint_bottleneck(prev_features)
        else:
            bottleneck_output = self.bottleneck_fn(prev_features)

        new_features = self.conv2(self.norm2(bottleneck_output))
        if self.drop_rate > 0:
            new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)
        return new_features


class DenseBlock(nn.ModuleDict):
    """DenseNet Block.

    Contains multiple dense layers with concatenated features.
    """
    _version = 2

    def __init__(

View on GitHub (pinned to 9a5261e31b)

Solutions

1. Disable gradient checkpointing before scripting: model.set_gradient_checkpointing(enable=False) or create the model without it
2. Use eager mode or torch.compile instead of torch.jit.script for checkpointed DenseNets
3. Script a separate model instance created purely for inference (no checkpointing, inputs with requires_grad=False)

Example fix

# before
model = timm.create_model('densenet121', grad_checkpointing=True)
scripted = torch.jit.script(model)
# after
model = timm.create_model('densenet121')  # no grad checkpointing
scripted = torch.jit.script(model)

if getattr(model, 'grad_checkpointing', False):
    model.set_gradient_checkpointing(enable=False)
scripted = torch.jit.script(model)

try:
    scripted = torch.jit.script(model)
except Exception as e:
    if 'Memory Efficient not supported in JIT' in str(e):
        model.set_gradient_checkpointing(enable=False)
        scripted = torch.jit.script(model)
    else:
        raise

When it happens

Trigger: Calling torch.jit.script(model) (or scripting a parent containing it) on a DenseNet created with gradient_checkpointing=True, followed by a forward pass where inputs require grad; also triggered by activation checkpointing utilities that run under torch.jit.is_scripting().

Common situations: Deploying a training-time memory-efficient DenseNet with TorchScript; wrapping timm models in a scripted pipeline while keep_grad flags leak through; toggling grad_checkpointing globally for a training utility then trying to export the same model.

Related errors

• Input image must have positive dimensions, got H={height}, W
• Invalid class map file, expected a dict ({class_map_path}).
• Dataset length is unknown, please pass `num_samples` explici
• Found 0 images in subfolders of {root}. Supported image exte
• Invalid or corrupt tar info cache file {cache_path}.

AI-assisted analysis of huggingface/pytorch-image-models@9a5261e31b (2026-08-27). Data as JSON: /api/errors/88e54d2f7b91d002. Report an issue: GitHub.