diffusion_policy/model/common/shape_util.py

from typing import Dict, List, Tuple, Callable
import torch
import torch.nn as nn

def get_module_device(m: nn.Module):
    device = torch.device('cpu')
    try:
        param = next(iter(m.parameters()))
        device = param.device
    except StopIteration:
        pass
    return device

@torch.no_grad()
def get_output_shape(
        input_shape: Tuple[int],
        net: Callable[[torch.Tensor], torch.Tensor]
    ):  
        device = get_module_device(net)
        test_input = torch.zeros((1,)+tuple(input_shape), device=device)
        test_output = net(test_input)
        output_shape = tuple(test_output.shape[1:])
        return output_shape