I was wondering about how can I implement an environment purely on GPU. Say, if all the variables in the environment class are torch.Tensor, well they stay on GPU during run-time?
Take the following environment for example:
class ENV_GPU(object):
    def __init__(self, a=3):
        self.num = torch.zeros((a,), dtype=torch.int8)
    def step(action):
        self.num[action] += 1