Nothing happens after death to the atma which is merely a witness. It is the subtle body that is affected by death. Check Gita 2.17-2.22 for the answer to your question. For example, Gita 2.22 says,
'Just as a man gives up old garments and puts on new ones, so the
embodied self abandons decrepit bodies and assumes new ones.
What happens immediately after death? The answer is given in Gita 15.8 which says,
When he gets a new body or abandons an old one, the Jiva, the lord of
the body, moves, carrying them (the mind and the senses) with him, as
the wind carries smells from their seats (in flowers and the like).
What this means is that the physical body dies but the mind and the senses leave the physical body.Persons who have not attained moksha are part of the life-death-life.... cycle. They may also have to stay in heaven or hell for some time if they have done a particularly good thing or a heinous crime, respectively, which cannot be adequately rewarded or punished in the Earth plane. After the merit of their particular good karma or the demerit of their particular bad karma is exhausted, they will again be born into the Earth plane. The process is linked to the concept of residual karma. This residual karma idea and what happens to evil-doers is explained in Brahma Sutra.
But of others (i.e. those who have not performed sacrifices etc.) the
ascent is to the abode of Yama, and after having experienced (the
results of their evil works) the descent (to the earth again takes
place). On account of such a passage (for the evil-doer) being
declared by the Sruti.
(Brahma Sutra Sankara Bhasya of Sankaracharya 3.1.13)
The Sruti passage referred to in the above shloka is:
The Hereafter never reveals itself to a person devoid of
discrimination, heedless, and perplexed by the delusion of wealth.
"This world alone exists," he thinks," and there is no other." Again
and again he comes under my sway.
(Yama speaks in Katha Upanishad I.2.6)
It is clear that if you do bad deeds that you end up in Yama Loka. However, there are two perplexing questions here. Does the person who descends to earth from Yama loka (or Swarga loka) return with any Karma (i.e. is the Karma of the evil-doer after serving his term in Yama loka zero?)? What decides that a person will after death be immediately reborn in earth or return to earth after serving a time in Yama loka? These two questions are actually linked. The answer to this question is given in Brahma Sutra:
On the exhaustion of (good) work (the soul) with residual Karma
(descends to this earth), as is known from the Sruti amd Smriti, along
the path (it) went by (from here) and differently too.
(Brahma Sutra Sankara Bhasya 3.1.8)
The above Sutra says that what is exhausted in heaven is only that specific Karma which gave the soul a birth in lower heaven, but on the exhaustion of this Karma the remaining Karma, good and bad, brings it back to earth. One cannot do karma in either swarga loka or in naraka. While this Sutra talks of heaven and good deed, it also applies to Yama Loka and bad deeds. Let me explain the idea by a concrete example. Let us suppose that a dead person has a list of 10 karmas with goodness values (1, -3, 4, 5, -100, 4, 3, 2, 5, 8) where the minus sign stands for evil deeds. What Brahma Sutra Sankara Bhasya 3.1.8 is saying is that this person will go to Yama Loka only for that extremely evil deed of goodness value -100. After he has exhausted the bad effect due to the -100 Karma, he can then go back to earth with the residual karma (1, -3, 4,5,4,3,2,5,8). The idea is that only extremely evil or good deeds lead a man to Yama loka or to lower heavens. All minor good and bad karmas are served on the earth plane.The person who has attained moksha is free forever from the endless cycle of Samsara. The cycle of birth and death will only end after the attainment of moksha or liberation. The length of the process will depend on how quickly or slowly the person will attain moksha.