ValueError: Data type <class 'NoneType'> not supported

[dario-coscia]

Describe the bug
DataCondition not correctly batch, yielding ValueError: Data type <class 'NoneType'> not supported

To Reproduce

import torch

from pina import Trainer
from pina.solver import SingleSolverInterface
from pina.problem import AbstractProblem
from pina.condition import Condition, DataCondition
from pina.model import FeedForward


class mySupervised(SingleSolverInterface):
    accepted_conditions_types = DataCondition

    def optimization_cycle(self, batch):
        condition_loss = {}
        for condition_name, points in batch:
            condition_loss[condition_name] = self.loss_data(
                input=points["input"]
            )
        return condition_loss
    
    def loss_data(self, input):
        return input.mean()


class GenerativeProblem(AbstractProblem):

    conditions = {
        # "data" : DataCondition(input=torch.randn(10, 2)) # the same error
        "data": Condition(input=torch.randn(10, 2))
    }
    input_variables = output_variables = None


trainer = Trainer(
    solver=mySupervised(
        problem=GenerativeProblem(), model=FeedForward(2, 2)
    ),
    max_epochs=10,
)

trainer.train()

[dario-coscia]

Setting shuffle=False in the Trainer leads to the following error:

Traceback (most recent call last):
  File ".../PINA/pina/data/data_module.py", line 431, in <dictcomp>
    k: v[offset : offset + stage_len]
TypeError: 'NoneType' object is not subscriptable