I was talking about in general NNAPI usage but didn’t run the above example myself to check the GPU acceleration. I’m sorry for the confusion.
This link 24 might be more related to you.
It says the primary targets of the vulkan backed are mobile GPUs and pytorch should be built with a proper cmake option, USE_VULKAN.
It seems like the mobile_optimizer 16, torch.utils.mobile_optimizer.optimize_for_mobile already supports mobile GPU if built with vulkan enabled.