Avoid gradient tracking

[aditya0by0]

For inference, PyTorch Lightning automatically calls model.eval() and wraps the forward pass in torch.inference_mode().

When performing inference outside Lightning, using only model.eval() is not sufficient.

• model.eval() does not disable gradients. It only switches layers like dropout and batchnorm into evaluation mode.
• Gradients are still tracked by autograd unless explicitly disabled.

Disabling gradient tracking is important because it:

1. Reduces memory usage by avoiding storing intermediate activations for backpropagation.
2. Speeds up inference by skipping autograd bookkeeping.

For this purpose, torch.inference_mode() is recommended. It is newer, faster, and more restrictive than torch.no_grad(), making it ideal for inference.

References:

• Behavior of model.eval() in PyTorch
• Difference between torch.no_grad() and torch.inference_mode()

[aditya0by0]

There is more thing we could to try for faster inference: torch.compile()

Please refer below to all the links:
https://huggingface.co/docs/transformers/en/perf_torch_compile
https://docs.pytorch.org/tutorials/intermediate/torch_compile_full_example.html
https://docs.pytorch.org/tutorials/intermediate/torch_compile_tutorial.html

[aditya0by0]

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 10)

    def forward(self, x):
        return torch.relu(self.linear(x))

# Original model
model = MyModel()

# Compile for faster execution
model = torch.compile(model, mode="default")

x = torch.randn(32, 10)
y = model(x)  # runs faster

[sfluegel05]

This is an interesting point. I am wondering if chebifier is the right place to implement this. You already started a PR here ChEB-AI/python-chebai#135 which looks like it tries to provide what we need here - a common function for all the functionality surrounding inference predictions.

If you generalise the logic in ChEB-AI/python-chebai#135 a bit further, you should be able to put all the inference_mode and torch.compile calls in chebai and only call a generic chebai function here.

[aditya0by0]

closing this PR, as agreed on the suggestion that prediction logic should be handled in respective data module and model module. Hence gradients problem will be addressed there.

This is an interesting point. I am wondering if chebifier is the right place to implement this. You already started a PR here ChEB-AI/python-chebai#135 which looks like it tries to provide what we need here - a common function for all the functionality surrounding inference predictions.

If you generalise the logic in ChEB-AI/python-chebai#135 a bit further, you should be able to put all the inference_mode and torch.compile calls in chebai and only call a generic chebai function here.