class ForceWithParameters(pt.nn.Module):

    def __init__(self):
        super(ForceWithParameters, self).__init__()

    def forward(
        self, positions: Tensor, parameter1: Tensor, parameter2: Tensor
    ) -> Tensor:
        x2 = positions.pow(2).sum(dim=1)
        u_harmonic = ((parameter1 + parameter2**2) * x2).sum()
        return u_harmonic


def example():
    # Create a random cloud of particles.
    numParticles = 10
    system = mm.System()
    positions = np.random.rand(numParticles, 3)
    for _ in range(numParticles):
        system.addParticle(1.0)

    # Create a force
    pt_force = ForceWithParameters()
    model = pt.jit.script(pt_force)
    tforce = ot.TorchForce(model)
    # Add some parameters
    parameter1 = 1.0
    parameter2 = 1.0
    force.setOutputsForces(False)
    force.addGlobalParameter("parameter1", parameter1)
    force.addEnergyParameterDerivative("parameter1")
    force.addGlobalParameter("parameter2", parameter2)
    force.addEnergyParameterDerivative("parameter2")
    # Enable energy derivatives for the parameter
    system.addForce(force)
    # Compute the forces and energy.
    integ = mm.VerletIntegrator(1.0)
    platform = mm.Platform.getPlatformByName(platform)
    context = mm.Context(system, integ, platform)
    context.setPositions(positions)
    state = context.getState(
        getEnergy=True, getForces=True, getParameterDerivatives=True
    )
    # See if the energy and forces and the parameter derivative are correct.
    # The network defines a potential of the form E(r) = (parameter1 + parameter2**2)*|r|^2
    r2 = np.sum(positions * positions)
    expectedEnergy = (parameter1 + parameter2**2) * r2
    assert np.allclose(
        r2,
        state.getEnergyParameterDerivatives()["parameter1"],
    )
    assert np.allclose(
        2 * parameter2 * r2,
        state.getEnergyParameterDerivatives()["parameter2"],
    )