在 pytorch 中同时训练多个模型

Question

我有 k 个相同的 pytorch 模型，所有这些模型都可以放入单个 GPU 的内存中。我想使用 k 个不同的输入批次（每个模型一个）在这些模型上运行单个训练步骤。我想在一次调用中完成所有操作 - 即我不想一次按顺序加载它们，而是将 k 个输入批次视为单个超级模型的单个输入，生成 k 个损失的张量，和 k 个梯度的张量。是否有任何可用的抽象可以实现这一点？