By the way, what if I need to calculate some intermediate results during every step? How can I make sure that there is no date moved between CPU and GPU? Take the following for example,
    def forward(self, action):
        self.knapsack_num += self.action2shift[action]
        self.knapsack_num = torch.clamp(self.knapsack_num, max=self.knapsack_max)
        reward = torch.zeros((1,), dtype=torch.float)
        if action == self.num_food_types:
            if torch.equal(self.expected_num, self.knapsack_num + self.warehouse_num):
                reward = 100 * torch.ones((1,), dtype=torch.float)
            else:
                reward = -100 * torch.ones((1,), dtype=torch.float)
            return self.knapsack_num, reward, True
        else:
            return self.knapsack_num, reward, False