gatelfpytorchjson.classificationmodule module

class gatelfpytorchjson.classificationmodule.ClassificationModule(inputlayersinfo, hiddenlayersinfo, outputlayerinfo, featureconfig)

Bases: torch.nn.modules.module.Module

forward(batch)

on_cuda()

Returns true or false depending on if the module is on cuda or not. Unfortunately there is no API method in PyTorch for this so we get this from the first parameter of the model and cache it.

set_seed(seed)