Од септембра 1943. до маја 1944. боравио је са Врховним штабом на ослобођеној територији Босанске крајине, а након немачког десанта на Дрвар и краћег задржавања у Барију, прешао је на Вис. Половином октобра 1944, са делом Врховног штаба, пребацио на ослобођену територију западне Србије и потом ослобођени Београд, одакле је учествовао у командовању операцијама за даље ослобођење земље, као и разговорима са представницима савезничких мисија око војне помоћи. Крајем децембра 1944, предводио је делегацију НОВЈ, која је у посети Совјетском Савезу боравила до фебруара 1945. године. Реорганизацијом НОВЈ у Југословенску армију, 1. марта 1945. постао је први начелник Генералштаба ЈА и ову дужност обављао до септембра исте године. Марта 1946. налазио се у групи официра ЈА који су отишли на школовање у Москву. Након завршене Војне академије „Ворошилов”, јуна 1948. вратио се у Југославију.