I tried the loop
for i in range(len(mask)):
      prediction_total = mask[i].float().view(-1, 1) * prediction

because of dimension mismatch, I used .view, but the answer is not correct:
>>>print(prediction)
tensor([[3.0603e-02, 5.6551e-02, 9.2520e-02,  ..., 4.0616e-02, 1.6195e-01,
         4.0573e-02],
        [2.9348e-02, 1.0917e-02, 1.8414e-04,  ..., 3.0553e-06, 9.5730e-01,
         2.0963e-03],
        [5.0212e-02, 9.1508e-03, 1.2127e-03,  ..., 5.4165e-05, 9.3203e-01,
         5.9912e-03],
        ...,
        [1.9600e-02, 1.7821e-01, 2.8131e-03,  ..., 7.4272e-03, 2.5914e-02,
         7.5438e-01],
        [1.4370e-01, 4.3823e-02, 1.7975e-03,  ..., 2.7101e-03, 1.4415e-01,
         6.6198e-01],
        [1.7615e-02, 2.3825e-02, 6.6172e-04,  ..., 1.4958e-02, 2.7619e-02,
         9.1287e-01]], device='cuda:0')

>>>print(mask[0])
tensor([1, 0, 0,  ..., 0, 0, 0], device='cuda:0', dtype=torch.uint8)

for i=0 >>>print(prediction_total)
tensor([[0.0306, 0.0566, 0.0925,  ..., 0.0406, 0.1619, 0.0406],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        ...,
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000]],
       device='cuda:0')

prediction_total supposed to be nonzero just in two values, but it is nonzero for the whole first row.