最近在做一个简单的AI Agent demo，需要让LLM多次推理并调用工具，比如先让模型决定调用搜索，再根据搜索结果生成回答。我目前是用torch.no_grad()包住每次推理，但感觉代码很乱，而且想知道这样会不会影响梯度流？如果后续想对Agent的某些步骤做强化学习微调，是不是得把LLM的调用都包在torch.enable_grad()里？另外，有没有现成的框架或库能帮忙管理这种多步推理的计算图？自己手写循环感觉容易出错，求大佬指点。