{"record":{"id":"88e54d2f7b91d002","repo":"huggingface/pytorch-image-models","slug":"memory-efficient-not-supported-in-jit","errorCode":null,"errorMessage":"Memory Efficient not supported in JIT","messagePattern":"Memory Efficient not supported in JIT","errorType":"exception","errorClass":"Exception","httpStatus":null,"severity":"error","filePath":"timm/models/densenet.py","lineNumber":100,"sourceCode":"    # torchscript does not yet support *args, so we overload method\n    # allowing it to take either a List[Tensor] or single Tensor\n    def forward(self, x: Union[torch.Tensor, List[torch.Tensor]]) -> torch.Tensor:  # noqa: F811\n        \"\"\"Forward pass.\n\n        Args:\n            x: Input features (single tensor or list of tensors).\n\n        Returns:\n            New features to be concatenated.\n        \"\"\"\n        if isinstance(x, torch.Tensor):\n            prev_features = [x]\n        else:\n            prev_features = x\n\n        if self.grad_checkpointing and self.any_requires_grad(prev_features):\n            if torch.jit.is_scripting():\n                raise Exception(\"Memory Efficient not supported in JIT\")\n            bottleneck_output = self.call_checkpoint_bottleneck(prev_features)\n        else:\n            bottleneck_output = self.bottleneck_fn(prev_features)\n\n        new_features = self.conv2(self.norm2(bottleneck_output))\n        if self.drop_rate > 0:\n            new_features = F.dropout(new_features, p=self.drop_rate, training=self.training)\n        return new_features\n\n\nclass DenseBlock(nn.ModuleDict):\n    \"\"\"DenseNet Block.\n\n    Contains multiple dense layers with concatenated features.\n    \"\"\"\n    _version = 2\n\n    def __init__(","sourceCodeStart":82,"sourceCodeEnd":118,"githubUrl":"https://github.com/huggingface/pytorch-image-models/blob/9a5261e31b3b5128526eb2658333b4c0a54464ae/timm/models/densenet.py#L82-L118","documentation":"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.","triggerScenarios":"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().","commonSituations":"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.","solutions":["Disable gradient checkpointing before scripting: model.set_gradient_checkpointing(enable=False) or create the model without it","Use eager mode or torch.compile instead of torch.jit.script for checkpointed DenseNets","Script a separate model instance created purely for inference (no checkpointing, inputs with requires_grad=False)"],"exampleFix":"# before\nmodel = timm.create_model('densenet121', grad_checkpointing=True)\nscripted = torch.jit.script(model)\n# after\nmodel = timm.create_model('densenet121')  # no grad checkpointing\nscripted = torch.jit.script(model)","handlingStrategy":"validation","validationCode":"if getattr(model, 'grad_checkpointing', False):\n    model.set_gradient_checkpointing(enable=False)\nscripted = torch.jit.script(model)","typeGuard":null,"tryCatchPattern":"try:\n    scripted = torch.jit.script(model)\nexcept Exception as e:\n    if 'Memory Efficient not supported in JIT' in str(e):\n        model.set_gradient_checkpointing(enable=False)\n        scripted = torch.jit.script(model)\n    else:\n        raise","preventionTips":["Separate training (checkpointed, eager) and deployment (non-checkpointed, scripted) model instances","Always disable grad checkpointing before export","Prefer torch.compile over TorchScript for new deployments"],"tags":["timm","densenet","torchscript","gradient-checkpointing"],"backgroundTag":"torchscript-unsupported-feature","analyzedSha":"9a5261e31b3b5128526eb2658333b4c0a54464ae","analyzedAt":"2026-08-27T02:34:25.417Z","schemaVersion":2},"datasetVersion":"2026-08-27T03:17:27.898Z"}