我正在尝试创建一个自定义的 pytorch 层，该层为一个具有 16 帧且大小为 3,224,224 的视频创建形状为 16,3,224,224 的高斯热图。该功能有效，但无法处理批处理数据。对于一批 10 个视频，我目前正在使用 for 循环遍历每个视频。这需要很多时间，我想并行处理视频。我怎样才能使函数像其他 pytorch 模型一样接受批量输入。
我尝试尽可能多地使用火炬操作员，但由于批量输入，我仍然遇到张量维度错误