.

main.py

@@ -40,8 +40,8 @@ model = initialize_model("config_final_model.yaml", "yuntian-deng/computer-model
 model = model.to(device)
 #model = torch.compile(model)
 
-padding_image = torch.zeros(1, SCREEN_HEIGHT // 8, SCREEN_WIDTH // 8, 4)
-padding_image = (padding_image - DATA_NORMALIZATION['mean']) / DATA_NORMALIZATION['std']
+padding_image = torch.zeros(1, SCREEN_HEIGHT // 8, SCREEN_WIDTH // 8, 16)
+padding_image = (padding_image - DATA_NORMALIZATION['mean']).view(1, 1, 1, -1) / DATA_NORMALIZATION['std'].view(1, 1, 1, -1)
 padding_image = padding_image.to(device)
 
 # Valid keyboard inputs