Solved by ptrblck in post #4 
              
              
                .to() is implemented in nn.Module. 
If you derive your class from nn.Module and define step as forward it should work: 
class ENV_GPU(nn.Module):
    def __init__(self, a=3):
        super(ENV_GPU, self).__init__()
        num = torch.zeros((a,), dtype=torch.int8)
        self.register_buffer('num',…