fix multiple issues on GPU

Author: hx89

src/maxdiffusion/max_utils.py

@@ -62,6 +62,8 @@
 
 from google.cloud import storage
 
+libcudart = cdll.LoadLibrary("libcudart.so")
+
 FrozenDict = core.frozen_dict.FrozenDict
 
 
@@ -78,12 +80,18 @@ def l2norm_pytree(x):
 
 def activate_profiler(config):
   if jax.process_index() == 0 and config.enable_profiler:
-    jax.profiler.start_trace(config.tensorboard_dir)
+    if config.profiler == 'nsys':
+      libcudart.cudaProfilerStart()
+    else:
+      jax.profiler.start_trace(config.tensorboard_dir)
 
 
 def deactivate_profiler(config):
   if jax.process_index() == 0 and config.enable_profiler:
-    jax.profiler.stop_trace()
+    if config.profiler == 'nsys':
+      libcudart.cudaProfilerStop()
+    else:
+      jax.profiler.stop_trace()
 
 
 def initialize_summary_writer(config):