[BugFix] Adaptable non-blocking for mps and non cuda device in batched-envs#1900