fix(nyz): fix multi trainer test

Author: PaParaZz1

ding/framework/middleware/tests/test_trainer.py

@@ -1,7 +1,8 @@
 import pytest
 import random
-import torch
 import copy
+import torch
+import treetensor.torch as ttorch
 from unittest.mock import Mock, patch
 from ding.data.buffer import DequeBuffer
 from ding.framework import OnlineRLContext, task
@@ -10,6 +11,7 @@
 
 
 class MockPolicy(Mock):
+    _device = 'cpu'
     # MockPolicy class for train mode
     def forward(self, train_data, **kwargs):
         res = {
@@ -19,6 +21,7 @@ def forward(self, train_data, **kwargs):
 
 
 class MultiStepMockPolicy(Mock):
+    _device = 'cpu'
     # MockPolicy class for multi-step train mode
     def forward(self, train_data, **kwargs):
         res = [
@@ -34,7 +37,7 @@ def forward(self, train_data, **kwargs):
 
 def get_mock_train_input():
     data = {'obs': torch.rand(2, 2), 'next_obs': torch.rand(2, 2), 'reward': random.random(), 'info': {}}
-    return data
+    return ttorch.as_tensor(data)
 
 
 @pytest.mark.unittest

docker/Dockerfile.base

@@ -3,7 +3,7 @@ FROM pytorch/pytorch:1.7.1-cuda11.0-cudnn8-runtime as base
 WORKDIR /ding
 
 RUN apt update \
-    && apt install libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev swig curl git vim gcc \g++ make wget locales dnsutils zip -y \
+    && apt install libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev swig curl git vim gcc \g++ make wget locales dnsutils zip unzip -y \
     && apt clean \
     && rm -rf /var/cache/apt/* \
     && sed -i '/en_US.UTF-8/s/^# //g' /etc/locale.gen \
@@ -33,7 +33,7 @@ RUN apt-get update && \
     python3.8 python3-pip python3.8-dev
 
 RUN apt update \
-    && apt install libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev swig curl git vim gcc \g++ make wget locales dnsutils zip -y \
+    && apt install libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev swig curl git vim gcc \g++ make wget locales dnsutils zip unzip -y \
     && apt clean \
     && rm -rf /var/cache/apt/* \
     && sed -i '/en_US.UTF-8/s/^# //g' /etc/locale.gen \