RuntimeError in avg_pool2d() which indicates that argument 'input' must be Variable

Hi there! I am trying to use nn.AvgPool2d and I am getting the following error: RuntimeError: avg_pool2d(): argument 'input' (position 1) must be Variable, not torch.cuda.FloatTensor
The error is generated by the following code:

import torch
import torch.nn as nn

m = nn.AvgPool2d(3, stride= 2)
inputs = torch.randn(20, 16, 50, 32)
output = m(inputs)

The output is the following:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 357, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/pooling.py", line 506, in forward
    self.padding, self.ceil_mode, self.count_include_pad)
RuntimeError: avg_pool2d(): argument 'input' (position 1) must be Variable, not torch.FloatTensor

I am using a docker container (Dockerfile) based on ubuntu 16.04 with cuda9.0 and cudnn7. Python 3.6 installing PyTorch from http://download.pytorch.org/whl/cu90/torch-0.3.1-cp36-cp36m-linux_x86_64.whl. I am not sure if this can be related with PyTorch version. Should I update to PyTorch 0.4.0 ?. The code which produces the above error was extracted from the official PyTorch nn documentation. Thanks in advance for your help! I really don’t know what is the problem

Updating to PyTorch 0.4.0 solved this issue