0710-1050

context_unet.py

@@ -516,7 +516,7 @@ class ContextUnet(nn.Module):
 
     def forward(self, x, timesteps, y=None):
         hs = []
-        print("device of timesteps, self.model_channels:", timesteps.device, self.model_channels)
+        # print("device of timesteps, self.model_channels:", timesteps.device, self.model_channels)
         emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
         if y != None:
             text_outputs = self.token_embedding(y.float())

diffusion.py

@@ -549,12 +549,15 @@ if __name__ == "__main__":
 # torch.cuda.set_device(0)
 
 # %%
-print("torch.cuda.is_available() =", torch.cuda.is_available())
-print("torch.cuda.device_count() =", torch.cuda.device_count())
+print(
+    "torch.cuda.is_available() =", torch.cuda.is_available(), 
+"torch.cuda.device_count() =", torch.cuda.device_count(),
+"torch.cuda.is_initialized() =", torch.cuda.is_initialized()
+)
 # print(torch.cuda.__dir__())
 
 # %%
-print("torch.cuda.is_initialized() =", torch.cuda.is_initialized())
+# print("torch.cuda.is_initialized() =", torch.cuda.is_initialized())
 # print("torch.cuda.get_device_name() =", torch.cuda.get_device_name())
 print("torch.cuda.current_device() =", torch.cuda.current_device())
 # print("torch.cuda.get_device_capability() =", torch.cuda.get_device_capability())