Training a Large Graph Neural Network for Time Series Traffic Prediction Problem

Question

I am training a GNN (GCN+GRU) for traffic prediction problem using pygeometric, my project is like the example shown in the official website https://pytorch-geometric-temporal.readthedocs.io/en/latest/notes/introduction.html

(https://i.stack.imgur.com/wGnFj.png)

the training process need to iterate through each snapshots for each epoch, and the loss is aggregated after looping through all the snapshots. for example:....

for each epoch: 
   loss_list = []                                                      
   for each snapshot:                                       
     prediction = model()                                                                           
     loss_list.append(loss)  
   lost = avg(loss_list)
   loss.backward

but I am very confused when I am dealing with large graph and need to perform mini batch training by having a loader that generate subgraphs, because I dont know where I should put the backward loss?

Option A:

for each epoch:                                                      
   for each snapshot:  
      for each subgraph:                           
         prediction = model()                                                                           
         loss.backward

which means the loss will be calculated for each subgraphs for each snapshot

Option B:

for each epoch:
  for each subgraph:  
      loss_list = []
      for each snapshot:                           
         prediction = model()  
         loss_list.append(loss)                                                                         
      lost = avg(loss_list)
      loss.backward
                                                 
   

which means the loss will be aggregated first after looping through snapshots and the backward action will be performed for each subgraph (although this option sounds more reasonable, just that i might need to build a dataloader on my own which will generate subgraphs that contains many snapshot data )

could someone tell me which way makes more sense?

Thank you!