Log gpu model

Pitfall2-Example.ipynb

@@ -69,7 +69,7 @@
     "    data.x = torch.ones((g.number_of_nodes(),1))\n",
     "    \n",
     "    for u in colored_nodes:\n",
-    "        data.x[u,0]=0\n",
+    "        data.x[u,0]=2\n",
     "\n",
     "    return data"
    ]

main.py

@@ -1,3 +1,4 @@
+import subprocess
 from enum import Enum
 
 import mlflow
@@ -23,6 +24,13 @@ def main(experiment: Experiment = typer.Argument(..., help="Dataset to use"),
                                        help="Convolution class. Can be GCNConv or GraphConv"),
          ):
     mlflow.set_experiment(experiment.value)
+    try:
+        out = subprocess.Popen(['nvidia-smi', '-L'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+        stdout, stderr = out.communicate()
+        gpu_model = stdout.decode().strip()
+        mlflow.log_param('GPU', gpu_model)
+    except FileNotFoundError:
+        pass
     class_map = {
         Experiment.infection: Infection,
         Experiment.community: Community,