{"record":{"id":"16c9d98f61d41111","repo":"WZMIAOMIAO/deep-learning-for-image-processing","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":"pytorch_classification/Test8_densenet/model.py","lineNumber":67,"sourceCode":"\n        return False\n\n    @torch.jit.unused\n    def call_checkpoint_bottleneck(self, inputs: List[Tensor]) -> Tensor:\n        def closure(*inp):\n            return self.bn_function(inp)\n\n        return cp.checkpoint(closure, *inputs)\n\n    def forward(self, inputs: Tensor) -> Tensor:\n        if isinstance(inputs, Tensor):\n            prev_features = [inputs]\n        else:\n            prev_features = inputs\n\n        if self.memory_efficient and self.any_requires_grad(prev_features):\n            if torch.jit.is_scripting():\n                raise Exception(\"memory efficient not supported in JIT\")\n\n            bottleneck_output = self.call_checkpoint_bottleneck(prev_features)\n        else:\n            bottleneck_output = self.bn_function(prev_features)\n\n        new_features = self.conv2(self.relu2(self.norm2(bottleneck_output)))\n        if self.drop_rate > 0:\n            new_features = F.dropout(new_features,\n                                     p=self.drop_rate,\n                                     training=self.training)\n\n        return new_features\n\n\nclass _DenseBlock(nn.ModuleDict):\n    _version = 2\n\n    def __init__(self,","sourceCodeStart":49,"sourceCodeEnd":85,"githubUrl":"https://github.com/WZMIAOMIAO/deep-learning-for-image-processing/blob/1ec3fe6f374fc9969973a61f819de25658595afa/pytorch_classification/Test8_densenet/model.py#L49-L85","documentation":"Inside _DenseLayer.forward, when self.memory_efficient is enabled and at least one input to the layer requires grad, the layer uses gradient checkpointing for the bottleneck. Checkpointing (torch.utils.checkpoint) is incompatible with TorchScript JIT, so the code explicitly raises when torch.jit.is_scripting() is true. This prevents silently producing wrong gradients under JIT compilation.","triggerScenarios":"Running torch.jit.script (or a model that gets scripted) on a DenseNet built with memory_efficient=True while some prev_features tensors require grad, so forward hits the checkpointing branch during scripting.","commonSituations":"Users copy the official DenseNet definition and enable memory_efficient=True for training, then export or compile the model with TorchScript for deployment/ONNX export; scripts often forget to disable memory efficient mode before scripting.","solutions":["Set memory_efficient=False on the _DenseLayer (or construct the model without memory-efficient checkpointing) before calling torch.jit.script","Script the model under torch.jit.is_scripting()-aware branches, or trace with torch.jit.trace instead of script if compatible","Separate the training model (memory efficient) from the deployment model (plain bn_function path) and script only the latter"],"exampleFix":"// before\nmodel = densenet121(memory_efficient=True)\nscripted = torch.jit.script(model)  # raises\n// after\nmodel = densenet121(memory_efficient=False)\nscripted = torch.jit.script(model)","handlingStrategy":"fallback","validationCode":"if torch.jit.is_scripting() and model.memory_efficient:\n    model = rebuild_model_without_checkpointing(model)  # set memory_efficient=False\nscripted = torch.jit.script(model)","typeGuard":"def is_scripting_with_checkpointing(model) -> bool:\n    return torch.jit.is_scripting() and getattr(model, 'memory_efficient', False)","tryCatchPattern":"try:\n    scripted = torch.jit.script(model)\nexcept Exception as e:\n    if 'memory efficient not supported' in str(e):\n        model.memory_efficient = False\n        scripted = torch.jit.script(model)\n    else:\n        raise","preventionTips":["Keep two model configurations: memory_efficient=True for eager training, False for scripted/exported models","Check torch.jit.is_scripting() before enabling checkpointing paths","Smoke-test torch.jit.script on your model in CI before deployment"],"tags":["pytorch","torchscript","jit","gradient-checkpointing"],"backgroundTag":"jit-unsupported-feature","analyzedSha":"1ec3fe6f374fc9969973a61f819de25658595afa","analyzedAt":"2026-08-30T09:19:11.901Z","schemaVersion":2},"datasetVersion":"2026-08-30T13:17:10.514Z"}