Since you are using single channel images, you could swap the batch and the channel dimension and use a grouped convolution:
class FilterNet(nn.Module):

    def __init__(self, obs_size, hidden_size, kernel_size, activation_func='tanh'):
        super(FilterNet, self).__init__()
        self.kernel_size = kernel_size

        self.fc_obs1 = nn.Linear(obs_size, hidden_size)
        self.fc_obs2 = nn.Linear(hidden_size, hidden_size)
        self.fc_obs3 = nn.Linear(hidden_size, hidden_size)
        self.fc_obs4 = nn.Linear(hidden_size, kernel_size * kernel_size)

        self.act_func = getattr(torch, activation_func)

    def forward(self, images, tau, use_loop=False):
        hidden = self.act_func(self.fc_obs1(tau))
        hidden = self.act_func(self.fc_obs2(hidden))
        hidden = self.act_func(self.fc_obs3(hidden))
        kernel_weigths = self.fc_obs4(hidden)

        if use_loop:
            heatmaps = []
            for weights_i, image_i in zip(kernel_weigths.chunk(kernel_weigths.size(0), dim=0), images.chunk(images.size(0), dim=0)):
                kernel = weights_i.view(1, 1, self.kernel_size, self.kernel_size)
                heatmap = F.conv2d(image_i, kernel, padding=2)
                heatmap = heatmap.squeeze()
                heatmaps += [heatmap]
            heatmaps = torch.stack(heatmaps, 0)
        else:
            kernel_weigths = kernel_weigths.view(kernel_weigths.size(0), 1, self.kernel_size, self.kernel_size)
            images = images.permute(1, 0, 2, 3)
            heatmaps = F.conv2d(images, kernel_weigths, padding=2, groups=images.size(1))
            
        return heatmaps


model = FilterNet(10, 10, 3)
images = torch.randn(16, 1, 24, 24)
tau = torch.randn(16, 10)

output_loop = model(images, tau, use_loop=True)
output_grouped = model(images, tau, use_loop=False)

print((output_loop - output_grouped).max())
> tensor(2.3842e-07, grad_fn=<MaxBackward1>)

This will apply each kernel in a single input channel, so that in fact we are applying each kernel on a single sample from the batch.