.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', num)        
        
    def forward(self, action):
        self.num[action] += 1

model = ENV_GPU()
model.to('cuda')
model(0)
print(model.num)
> tensor([1, 0, 0], device='cuda:0', dtype=torch.int8)