В сентябре 1944 года гвардии генерал-майор Козин был вновь назначен на должность командира 52-й гвардейской стрелковой дивизии 3-й ударной армии 2-го Прибалтийского фронта, которой командовал до конца войны. Принимал участие в Прибалтийской наступательной операции, в ходе которой дивизия одной из первых вошла в Ригу. В декабре дивизию вместе с армией передали на 1-й Белорусский фронт, где она вновь отлично действовала в Висло-Одерской и Восточно-Померанской наступательных операциях. В ходе последней операции дивизия освободила около 90 населённых пунктов, в том числе 8 городов.