ErrorFeedback.py

from FunctionErrorFeedback import ErrorFeedbackFunction
from torch.autograd import Function, Variable
import torch
import torch.nn as nn

class EF(nn.Module):
    
    def __init__(self, layer_dim, error_dim):
        super(EF, self).__init__()
        self.feedback = torch.Tensor(error_dim, layer_dim)
        self.reset_variables()
    
    def reset_variables(self):
        # brutal initialization
        self.feedback.uniform_(-0.1, 0.1)
                
    def forward(self, input):
        return ErrorFeedbackFunction(input, self.feedback)