请教大佬代码实现问题

[LUO77123]

大佬您好，很感谢您的工作，我想将输入size(B , N, 5)修改为(B N, 5)，其他代码都能改，但是在cuda_op里面的不知道怎么改，例如cuda_ext.py中，我尝试修改v,m,nv的维度，将B , N修改成BN，但是输出不对，所以想请问大佬这块如何修改喃，谢谢

[lilanxiao]

你在最外面加一层Wrapper可不可以呢？这样就不用改任何代码了。例如这样：

from torch import Tensor
from oriented_iou_loss import cal_giou

def cal_giou_wrapper(box1: Tensor, box2:Tensor, bs: int):
    """
    box1 and box2: with the shape [B*N, 5]
    bs: give a batch size to make cal_giou happy
    """
    box1 = box1.view(bs, -1, 5)
    box2 = box2.view(bs, -1, 5)
    loss, iou = cal_giou(box1, box2)
    loss = loss.view(-1)
    iou = iou.view(-1)
    return loss, iou

如果你实在想自己改的话，你可以看看debug branch的版本。那个版本没有CUDA算子，应改要容易改一些。

[LUO77123]

你在最外面加一层Wrapper可不可以呢？这样就不用改任何代码了。例如这样：

from torch import Tensor
from oriented_iou_loss import cal_giou

def cal_giou_wrapper(box1: Tensor, box2:Tensor, bs: int):
    """
    box1 and box2: with the shape [B*N, 5]
    bs: give a batch size to make cal_giou happy
    """
    box1 = box1.view(bs, -1, 5)
    box2 = box2.view(bs, -1, 5)
    loss, iou = cal_giou(box1, box2)
    loss = loss.view(-1)
    iou = iou.view(-1)
    return loss, iou

如果你实在想自己改的话，你可以看看debug branch的版本。那个版本没有CUDA算子，应改要容易改一些。

谢谢大佬回复，我后面改不了了，就直接在加了一个维度，计算后在去掉，就不用修改程序了