fix rng generator to work in cpu-only environment too

models/switti.py

@@ -71,7 +71,7 @@ class Switti(nn.Module):
         self.rope = rope
 
         self.num_stages_minus_1 = len(self.patch_nums) - 1
-        self.rng = torch.Generator(device="cuda")
+        self.rng = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
 
         # 1. input (word) embedding
         self.word_embed = nn.Linear(self.Cvae, self.C)