Update test.py

test.py

@@ -9,7 +9,7 @@ setattr(torch.nn.Linear, 'reset_parameters', lambda self: None)     # 禁用默
 setattr(torch.nn.LayerNorm, 'reset_parameters', lambda self: None)  # 禁用默认参数init以获得更快的速度
 from models import VQVAE, build_vae_var
 
-os.environ["CUDA_VISIBLE_DEVICES"] = "2"
+os.environ["HIP_VISIBLE_DEVICES"] = "2"
 print(torch.cuda.get_device_name(0))
 MODEL_DEPTH = 16    # TODO:更改此处，指定模型
 assert MODEL_DEPTH in {16, 20, 24, 30}