class_docstr: The base script for running different 3DGN methods.
class_docstr_tok: ["The","base","script","for","running","different","3DGN","methods","."]
docstr: The script for training.
  
  Args:
      model: Which 3DGN model to use. Should be one of the SchNet, DimeNetPP, and SphereNet.
      optimizer (Optimizer): Pytorch optimizer for trainable parameters in training.
      train_loader (Dataloader): Dataloader for training.
      energy_and_force (bool, optional): If set to :obj:`True`, will predict energy and take the minus derivative of the energy with respect to the atomic positions as predicted forces. (default: :obj:`False`)    
      p (int, optinal): The forces’ weight for a joint loss of forces and conserved energy during training. (default: :obj:`100`)
      loss_func (function): The used loss funtion for training. 
      device (torch.device): The device where the model is deployed.
  
  :rtype: Traning loss. ( :obj:`mae`)
docstr_tok: ["The","script","for","training",".","Args",":","model",":","Which","3DGN","model","to","use",".","Should","be","one","of","the","SchNet",",","DimeNetPP",",","and","SphereNet",".","optimizer","(","Optimizer",")",":","Pytorch","optimizer","for","trainable","parameters","in","training",".","train_loader","(","Dataloader",")",":","Dataloader","for","training",".","energy_and_force","(","bool",",","optional",")",":","If","set","to",":","obj",":","`","True","`",",","will","predict","energy","and","take","the","minus","derivative","of","the","energy","with","respect","to","the","atomic","positions","as","predicted","forces",".","(","default",":",":","obj",":","`","False","`",")","p","(","int",",","optinal",")",":","The","forces","\u2019","weight","for","a","joint","loss","of","forces","and","conserved","energy","during","training",".","(","default",":",":","obj",":","`","100","`",")","loss_func","(","function",")",":","The","used","loss","funtion","for","training",".","device","(","torch.device",")",":","The","device","where","the","model","is","deployed",".",":","rtype",":","Traning","loss",".","(",":","obj",":","`","mae","`",")"]
code: def train(
      self,
      model,
      optimizer,
      train_loader,
      energy_and_force,
      p,
      loss_func,
      device,
  ):
      r"""
      The script for training.
  
      Args:
          model: Which 3DGN model to use. Should be one of the SchNet, DimeNetPP, and SphereNet.
          optimizer (Optimizer): Pytorch optimizer for trainable parameters in training.
          train_loader (Dataloader): Dataloader for training.
          energy_and_force (bool, optional): If set to :obj:`True`, will predict energy and take the minus derivative of the energy with respect to the atomic positions as predicted forces. (default: :obj:`False`)
          p (int, optinal): The forces’ weight for a joint loss of forces and conserved energy during training. (default: :obj:`100`)
          loss_func (function): The used loss funtion for training.
          device (torch.device): The device where the model is deployed.
  
      :rtype: Traning loss. ( :obj:`mae`)
  
      """
      model.train()
      loss_accum = 0
      for step, batch_data in enumerate(tqdm(train_loader)):
          optimizer.zero_grad()
          batch_data = batch_data.to(device)
          out = model(batch_data)
          if energy_and_force:
              force = -grad(
                  outputs=out,
                  inputs=batch_data.pos,
                  grad_outputs=torch.ones_like(out),
                  create_graph=True,
                  retain_graph=True,
              )[0]
              e_loss = loss_func(out, batch_data.y.unsqueeze(1))
              f_loss = loss_func(force, batch_data.force)
              loss = e_loss + p * f_loss
          else:
              loss = loss_func(out, batch_data.y.unsqueeze(1))
          loss.backward()
          optimizer.step()
          loss_accum += loss.detach().cpu().item()
      return loss_accum / (step + 1)
  
file_path: DIG/dig/threedgraph/method/run.py
imports: import time
  import os
  import torch
  from torch.optim import Adam
  from torch_geometric.data import DataLoader
  import numpy
  from torch.autograd import grad
  from torch.utils.tensorboard import SummaryWriter
  from torch.optim.lr_scheduler import StepLR
  from tqdm import tqdm
label_desc: Use image node_id 3 for calling the run obj's underlying member method code with example usage: obj.train(model, optimizer, train_loader, energy_and_force, p, loss_func, device) and returns: unknown