我想实现在单个服务器上的多个GPU上运行我的pytorch模型训练代码。
具体场景如下：训练时期=2000，每个时期的训练数据集的总数=1000，有三个GPU。在一个时期中每个GPU的训练数据运行的数量是episode//3。但由于显存的限制，每个GPU的显存容量只能运行批=32的数据。请问如何设置数据加载和数据加载器打包以适应上述情况，并在此过程中记录每次GPU运行的episode//3数据的平均值？
请注意，我的损失的设置是每个时期中每个事件的训练数据目标值的平均值。
目前在GPU上运行的结构如下：

for epoch in epochs:
    Avg_loss = 0
    while episode < train_num_episode:
          remaining = train_num_episode - episode
          batch_size = min(self.trainer_params['train_batch_size'], remaining)
          avg_loss = train_one_batch(batch_size)
    Avg_loss = average.update(avg_loss)