Why is torch complaining about an empty tensor?

The structure_loss method is supposed to return a loss for ground truth vs predicted masks:

import numpy as np
import torch
import torch.nn.functional as F
import torch.nn as nn

def structure_loss(pred, mask):
    weit = 1 + 5 * torch.abs(nn.functional.avg_pool2d(mask, kernel_size=31, stride=1, padding=15) - mask) # here
    wbce = F.binary_cross_entropy_with_logits(pred, mask, reduce='none')
    wbce = (weit * wbce).sum(dim=(2, 3)) / weit.sum(dim=(2, 3))

    pred = torch.sigmoid(pred)
    inter = ((pred * mask) * weit).sum(dim=(2, 3))
    union = ((pred + mask) * weit).sum(dim=(2, 3))
    wiou = 1 - (inter + 1) / (union - inter + 1)
    return (wbce + wiou).mean()


a = torch.Tensor(np.random.uniform(0,1, (352, 352)))
b = torch.Tensor(np.random.uniform(0,1, (352, 352)))
print(structure_loss(a,b))

However, I test it on dummy data I get

 File "binary_cross_entropy_for_random_array.py", line 8, in structure_loss
    weit = 1 + 5 * torch.abs(nn.functional.avg_pool2d(mask, kernel_size=31, 
stride=1, padding=15) - mask) # here
RuntimeError: non-empty 2D or 3D (batch mode) tensor expected for input

Process finished with exit code 1

How to generate tensors suitable for this method?