После присоединения Эстонии к СССР в 1940 году направлен в Эстонию, заведующий отделом Центрального комитета КП(б) Эстонии в Таллине. В августе 1941 года эвакуировался в Ленинград, откуда в декабре был вывезен через Ладожское озеро и отправлен в Челябинск. С 1942 года в составе 8-го стрелкового корпуса (эстонского), некоторое время комиссар 7-й стрелковой дивизии — согласно собственным воспоминаниям Штейна, убедил главнокомандующего Иосифа Сталина направить сформированные из эстонцев части не на Сталинградскую битву, а на освобождение Эстонии.