some horrible state but I need to stop developing on perlmutter

examples/demo.py

@@ -35,6 +35,8 @@ def demo_basic():
         print(i)
         optimizer.zero_grad()
         outputs = ddp_model(torch.randn(20, 10))
+        mat = torch.rand([3,3], dtype=torch.complex64)
+        inv_mat = torch.inverse(mat)
         labels = torch.randn(20, 5).to(device_id)
         loss_fn(outputs, labels).backward()
         optimizer.step()

examples/multi_gpu_simple_ptycho.py

@@ -34,7 +34,7 @@ def demo():
     sampler = torchdata.distributed.DistributedSampler(dataset)
     # Make a dataloader
     #data_loader = torchdata.DataLoader(dataset, batch_size=20,
-    #                                   shuffle=True)
+    #                                   shuffle=True, num_workers=0)
     data_loader = torchdata.DataLoader(dataset, batch_size=20,
                                        shuffle=False, sampler=sampler)
     
@@ -42,26 +42,29 @@ def demo():
     dataset.get_as(device=device)
     
     # Define the optimizer
-    optimizer = MyAdam(model.parameters(), lr=0.01)
+    #optimizer = MyAdam(model.parameters(), lr=0.01)
+    optimizer = t.optim.Adam(model.parameters(), lr=0.01)
     
     normalization=0
-    for inputs, patterns in data_loader:
+    for inputs, patterns in dataset:#data_loader:
         normalization += t.sum(patterns)
         
     dist.all_reduce(normalization)
 
-    for it in range(100):
+    for it in range(10):
+        print(it)
         sampler.set_epoch(it)
         loss = 0
         N = 0
         t0 = time.time()
         for inputs, patterns in data_loader:
+            #print(inputs)
             N += 1
             def closure():
                 optimizer.zero_grad()
                 
                 sim_patterns = model.forward(*inputs)
-            
+                
                 if hasattr(model, 'mask'):
                     loss = model.loss(patterns,sim_patterns, mask=model.mask)
                 else:
@@ -70,12 +73,11 @@ def demo():
                 loss.backward()
                 return loss.detach()
 
-            loss += optimizer.step(closure).detach()#.cpu().numpy()
+            loss += optimizer.step(closure)#.cpu().numpy()
         #dist.all_reduce(loss)
         #loss = dist.all_reduce(loss)
         #print(it, 'time', time.time()-t0)
         #print(loss / normalization)
-        
     return loss#.cpu().numpy()
 
 if __name__=='__main__':

examples/simple_ptycho.py

@@ -18,7 +18,7 @@ model = cdtools.models.SimplePtycho.from_dataset(dataset)
 #    def __getattr__(self, name):
 #        return getattr(self.module, name)
 
-model = t.nn.DataParallel(model, device_ids=[0,1,2,3])
+#model = t.nn.DataParallel(model, device_ids=[0,1,2,3])
 
 
 # Make a dataloader